// geneticPredictionModels provides the data structures and functions to represent, encode, and decode genetic prediction models
// Prediction models are used to predict polygenic disease risk scores and trait outcomes

package geneticPredictionModels

import "gorgonia.org/gorgonia"
import "gorgonia.org/tensor"

import "bytes"
import "encoding/gob"
import "errors"


type NeuralNetwork struct{

	// ExprGraph is a data structure for a directed acyclic graph (of expressions).
	Graph *gorgonia.ExprGraph

	// These are the weights for each layer of neurons
	Weights1 *gorgonia.Node
	Weights2 *gorgonia.Node
	Weights3 *gorgonia.Node

	// This is the computed prediction
	Prediction *gorgonia.Node
}

// This function returns the weights of the neural network
// We need this for training
func (inputNetwork *NeuralNetwork)GetLearnables()gorgonia.Nodes{

	weights1 := inputNetwork.Weights1
	weights2 := inputNetwork.Weights2
	weights3 := inputNetwork.Weights3

	result := gorgonia.Nodes{weights1, weights2, weights3}

	return result
}

// We use this to store a neural network's weights as a .gob file
type neuralNetworkForEncoding struct{

	// These are the weights for each layer of neurons
	Weights1 []float32
	Weights2 []float32
	Weights3 []float32

	// These represent the quantity of rows and columns for each weight layer
	Weights1Rows int
	Weights1Columns int
	Weights2Rows int
	Weights2Columns int
	Weights3Rows int
	Weights3Columns int
}

func EncodeNeuralNetworkObjectToBytes(inputNeuralNetwork NeuralNetwork)([]byte, error){

	weights1 := inputNeuralNetwork.Weights1
	weights2 := inputNeuralNetwork.Weights2
	weights3 := inputNeuralNetwork.Weights3

	weights1Slice := weights1.Value().Data().([]float32)
	weights2Slice := weights2.Value().Data().([]float32)
	weights3Slice := weights3.Value().Data().([]float32)

	weights1Rows := weights1.Shape()[0]
	weights1Columns := weights1.Shape()[1]
	weights2Rows := weights2.Shape()[0]
	weights2Columns := weights2.Shape()[1]
	weights3Rows := weights3.Shape()[0]
	weights3Columns := weights3.Shape()[1]

	newNeuralNetworkForEncoding := neuralNetworkForEncoding{
		Weights1: weights1Slice,
		Weights2: weights2Slice,
		Weights3: weights3Slice,

		Weights1Rows: weights1Rows,
		Weights1Columns: weights1Columns,
		Weights2Rows: weights2Rows,
		Weights2Columns: weights2Columns,
		Weights3Rows: weights3Rows,
		Weights3Columns: weights3Columns,
	}

	buffer := new(bytes.Buffer)

	encoder := gob.NewEncoder(buffer)

	err := encoder.Encode(newNeuralNetworkForEncoding)
	if (err != nil) { return nil, err }

	neuralNetworkBytes := buffer.Bytes()

	return neuralNetworkBytes, nil
}

func DecodeBytesToNeuralNetworkObject(inputNeuralNetwork []byte)(NeuralNetwork, error){

	if (inputNeuralNetwork == nil){
		return NeuralNetwork{}, errors.New("DecodeBytesToNeuralNetworkObject called with nil inputNeuralNetwork.")
	}

	buffer := bytes.NewBuffer(inputNeuralNetwork)

	decoder := gob.NewDecoder(buffer)

	var newNeuralNetworkForEncoding neuralNetworkForEncoding

	err := decoder.Decode(&newNeuralNetworkForEncoding)
	if (err != nil){ return NeuralNetwork{}, err }

	weights1 := newNeuralNetworkForEncoding.Weights1
	weights2 := newNeuralNetworkForEncoding.Weights2
	weights3 := newNeuralNetworkForEncoding.Weights3

	weights1Rows := newNeuralNetworkForEncoding.Weights1Rows
	weights1Columns := newNeuralNetworkForEncoding.Weights1Columns
	weights2Rows := newNeuralNetworkForEncoding.Weights2Rows
	weights2Columns := newNeuralNetworkForEncoding.Weights2Columns
	weights3Rows := newNeuralNetworkForEncoding.Weights3Rows
	weights3Columns := newNeuralNetworkForEncoding.Weights3Columns

	// This is the graph object we add each layer to
	newGraph := gorgonia.NewGraph()

	// A layer is a column of neurons
	// Each neuron has an initial value between 0 and 1
	getNewNeuralNetworkLayerWeights := func(layerName string, layerNeuronRows int, layerNeuronColumns int, layerWeightsList []float32)*gorgonia.Node{

		layerNameObject := gorgonia.WithName(layerName)

		layerBacking := tensor.WithBacking(layerWeightsList)
		layerShape := tensor.WithShape(layerNeuronRows, layerNeuronColumns)
		layerTensor := tensor.New(layerBacking, layerShape)

		layerValueObject := gorgonia.WithValue(layerTensor)

		layerObject := gorgonia.NewMatrix(newGraph, tensor.Float32, layerNameObject, layerValueObject)

		return layerObject
	}

	layer1 := getNewNeuralNetworkLayerWeights("Weights1", weights1Rows, weights1Columns, weights1)
	layer2 := getNewNeuralNetworkLayerWeights("Weights2", weights2Rows, weights2Columns, weights2)
	layer3 := getNewNeuralNetworkLayerWeights("Weights3", weights3Rows, weights3Columns, weights3)

	newNeuralNetworkObject := NeuralNetwork{

		Graph: newGraph,

		Weights1: layer1,
		Weights2: layer2,
		Weights3: layer3,
	}

	return newNeuralNetworkObject, nil
}


// This function will take a neural network and input layer and build the network to be able to compute a prediction
// We need to run a virtual machine after calling this function in order for the prediction to be generated
func (inputNetwork *NeuralNetwork)BuildNeuralNetwork(inputLayer *gorgonia.Node, predictionIsNumeric bool)error{

	// We copy node pointer (says to do this in a resource i'm reading)

	inputLayerCopy := inputLayer

	// We multiply weights at each layer and perform ReLU (Rectification) after each multiplication

	weights1 := inputNetwork.Weights1

	layer1Product, err := gorgonia.Mul(inputLayerCopy, weights1)
	if (err != nil) {
		return errors.New("Layer 1 multiplication failed: " + err.Error())
	}

	layer1ProductRectified, err := gorgonia.Rectify(layer1Product)
	if (err != nil){
		return errors.New("Layer 1 Rectify failed: " + err.Error())
	}

	weights2 := inputNetwork.Weights2

	layer2Product, err := gorgonia.Mul(layer1ProductRectified, weights2)
	if (err != nil) {
		return errors.New("Layer 2 multiplication failed: " + err.Error())
	}

	layer2ProductRectified, err := gorgonia.Rectify(layer2Product)
	if (err != nil){
		return errors.New("Layer 2 Rectify failed: " + err.Error())
	}

	weights3 := inputNetwork.Weights3

	layer3Product, err := gorgonia.Mul(layer2ProductRectified, weights3)
	if (err != nil) {
		return errors.New("Layer 3 multiplication failed: " + err.Error())
	}

	if (predictionIsNumeric == false){

		// We SoftMax the output to get the prediction

		prediction, err := gorgonia.SoftMax(layer3Product)
		if (err != nil) {
			return errors.New("SoftMax failed: " + err.Error())
		}

		inputNetwork.Prediction = prediction

	} else {

		// We Sigmoid the output to get the prediction

		prediction, err := gorgonia.Sigmoid(layer3Product)
		if (err != nil) {
			return errors.New("Sigmoid failed: " + err.Error())
		}

		inputNetwork.Prediction = prediction
	}

	return nil
}