/*
 * MIT License
 *
 * Copyright (c) 2019 Alexey Edelev <semlanik@gmail.com>
 *
 * This file is part of NeuralNetwork project https://git.semlanik.org/semlanik/NeuralNetwork
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this
 * software and associated documentation files (the "Software"), to deal in the Software
 * without restriction, including without limitation the rights to use, copy, modify,
 * merge, publish, distribute, sublicense, and/or sell copies of the Software, and
 * to permit persons to whom the Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be included in all copies
 * or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
 * PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
 * FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

package remotecontrol

import (
	context "context"
	fmt "fmt"
	"log"
	"net"

	"google.golang.org/grpc/codes"
	status "google.golang.org/grpc/status"

	neuralnetworkbase "../neuralnetworkbase"
	"gonum.org/v1/gonum/mat"
	grpc "google.golang.org/grpc"
)

type RemoteControl struct {
	nn               *neuralnetworkbase.NeuralNetwork
	activationsQueue chan *LayerMatrix
	biasesQueue      chan *LayerMatrix
	weightsQueue     chan *LayerMatrix
}

func (rw *RemoteControl) Init(nn *neuralnetworkbase.NeuralNetwork) {
	rw.nn = nn
	rw.activationsQueue = make(chan *LayerMatrix, 10)
	rw.biasesQueue = make(chan *LayerMatrix, 10)
	rw.weightsQueue = make(chan *LayerMatrix, 10)
}

func (rw *RemoteControl) UpdateActivations(l int, a *mat.Dense) {
	matrix := NewLayerMatrix(l, a, LayerMatrix_Activations)
	select {
	case rw.activationsQueue <- matrix:
	default:
	}
}

func (rw *RemoteControl) UpdateBiases(l int, biases *mat.Dense) {
	matrix := NewLayerMatrix(l, biases, LayerMatrix_Biases)
	select {
	case rw.biasesQueue <- matrix:
	default:
	}
}

func (rw *RemoteControl) UpdateWeights(l int, weights *mat.Dense) {
	matrix := NewLayerMatrix(l, weights, LayerMatrix_Weights)
	select {
	case rw.weightsQueue <- matrix:
	default:
	}
}

func NewLayerMatrix(l int, dense *mat.Dense, contentType LayerMatrix_ContentType) (matrix *LayerMatrix) {
	buffer, err := dense.MarshalBinary()
	if err != nil {
		log.Fatalln("Invalid dense is provided for remote control")
	}

	matrix = &LayerMatrix{
		Matrix: &Matrix{
			Matrix: buffer,
		},
		Layer:       int32(l),
		ContentType: contentType,
	}

	return
}

func (rw *RemoteControl) GetConfiguration(context.Context, *None) (*Configuration, error) {
	config := &Configuration{}

	for _, size := range rw.nn.Sizes {
		config.Sizes = append(config.Sizes, int32(size))
	}
	return config, nil
}

func (rw *RemoteControl) Activations(_ *None, srv RemoteControl_ActivationsServer) error {
	ctx := srv.Context()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		msg := <-rw.activationsQueue
		srv.Send(msg)
	}
}

func (rw *RemoteControl) Biases(_ *None, srv RemoteControl_BiasesServer) error {
	ctx := srv.Context()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		msg := <-rw.biasesQueue
		srv.Send(msg)
	}
}

func (rw *RemoteControl) Weights(_ *None, srv RemoteControl_WeightsServer) error {
	ctx := srv.Context()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		msg := <-rw.weightsQueue
		srv.Send(msg)
	}
}

func (rw *RemoteControl) Predict(context.Context, *Matrix) (*Matrix, error) {
	return nil, status.Error(codes.Unimplemented, "Not implemented")
}

func (rw *RemoteControl) Run() {
	grpcServer := grpc.NewServer()
	RegisterRemoteControlServer(grpcServer, rw)
	lis, err := net.Listen("tcp", "localhost:65001")
	if err != nil {
		fmt.Printf("Failed to listen: %v\n", err)
	}

	fmt.Printf("Listen localhost:65001\n")
	if err := grpcServer.Serve(lis); err != nil {
		fmt.Printf("Failed to serve: %v\n", err)
	}
}