/*
 * 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"
	"os"
	"sync"
	"time"

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

	neuralnetwork "../neuralnetwork"
	"gonum.org/v1/gonum/mat"
	grpc "google.golang.org/grpc"

	training "../training"
)

type RemoteControl struct {
	nn               *neuralnetwork.NeuralNetwork
	activationsQueue chan *LayerMatrix
	biasesQueue      chan *LayerMatrix
	weightsQueue     chan *LayerMatrix
	stateQueue       chan int
	mutex            sync.Mutex
	config           *Configuration
}

func (rw *RemoteControl) Init(nn *neuralnetwork.NeuralNetwork) {
	rw.nn = nn
	rw.activationsQueue = make(chan *LayerMatrix, 5)
	rw.biasesQueue = make(chan *LayerMatrix, 5)
	rw.weightsQueue = make(chan *LayerMatrix, 5)
	rw.stateQueue = make(chan int, 2)
	rw.config = &Configuration{}
	for _, size := range rw.nn.Sizes {
		rw.config.Sizes = append(rw.config.Sizes, int32(size))
	}
}

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 (rw *RemoteControl) UpdateState(state int) {
	select {
	case rw.stateQueue <- state:
	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) {
	return rw.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
		fmt.Println("Send Activations")
		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
		fmt.Println("Send Biases")
		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
		fmt.Println("Send Weights")
		srv.Send(msg)
	}
}

func (rw *RemoteControl) State(_ *None, srv RemoteControl_StateServer) error {
	ctx := srv.Context()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		state := <-rw.stateQueue
		msg := &NetworkState{
			State: NetworkState_State(state),
		}
		fmt.Printf("Send state %v %v\n", msg, state)
		srv.Send(msg)
	}
}

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

func (rw *RemoteControl) DummyStart(context.Context, *None) (*None, error) {
	go func() {
		rw.mutex.Lock()
		defer rw.mutex.Unlock()
		// trainer := training.NewMNISTReader("./minst.data", "./mnist.labels")
		trainer := training.NewTextDataReader("wine.data", 5)
		rw.nn.Train(trainer, 500)

		// for i := 0; i < nn.Count; i++ {
		// 	if i > 0 {
		// 		fmt.Printf("Weights after:\n%v\n\n", mat.Formatted(nn.Weights[i], mat.Prefix(""), mat.Excerpt(0)))
		// 		fmt.Printf("Biases after:\n%v\n\n", mat.Formatted(nn.Biases[i], mat.Prefix(""), mat.Excerpt(0)))
		// 		fmt.Printf("Z after:\n%v\n\n", mat.Formatted(nn.Z[i], mat.Prefix(""), mat.Excerpt(0)))
		// 	}
		// 	fmt.Printf("A after:\n%v\n\n", mat.Formatted(nn.A[i], mat.Prefix(""), mat.Excerpt(0)))
		// }

		outFile, err := os.OpenFile("./data", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0666)
		if err != nil {
			log.Fatal(err)
		}
		defer outFile.Close()
		rw.nn.SaveState(outFile)
		outFile.Close()

		rw.UpdateState(neuralnetwork.StateLearning)
		defer rw.UpdateState(neuralnetwork.StateIdle)
		failCount := 0
		trainer.Reset()
		for trainer.NextValidator() {
			dataSet, expect := trainer.GetValidator()
			index, _ := rw.nn.Predict(dataSet)
			//TODO: remove this is not used for visualization
			time.Sleep(400 * time.Millisecond)
			if expect.At(index, 0) != 1.0 {
				failCount++
				// fmt.Printf("Fail: %v, %v\n\n", trainer.ValidationIndex(), expect.At(index, 0))
			}
			if !trainer.NextValidator() {
				fmt.Printf("Fail count: %v\n\n", failCount)
				failCount = 0
				trainer.Reset()
			}
		}

		fmt.Printf("Fail count: %v\n\n", failCount)
		failCount = 0
		trainer.Reset()
		rw.UpdateState(neuralnetwork.StateIdle)
	}()

	return &None{}, nil
}

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)
	}
}