package models

import (
	"errors"
	"os"
	"path"

	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/db_types"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/tasks/utils"

	tf "github.com/galeone/tensorflow/tensorflow/go"
	"github.com/galeone/tensorflow/tensorflow/go/op"
	tg "github.com/galeone/tfgo"
	"github.com/galeone/tfgo/image"
)

func ReadPNG(scope *op.Scope, imagePath string, channels int64) *image.Image {
	scope = tg.NewScope(scope)
	contents := op.ReadFile(scope.SubScope("ReadFile"), op.Const(scope.SubScope("filename"), imagePath))
	output := op.DecodePng(scope.SubScope("DecodePng"), contents, op.DecodePngChannels(channels))
	output = op.ExpandDims(scope.SubScope("ExpandDims"), output, op.Const(scope.SubScope("axis"), []int32{0}))
	image := &image.Image{
		Tensor: tg.NewTensor(scope, output)}
	return image.Scale(0, 255)
}

func ReadJPG(scope *op.Scope, imagePath string, channels int64) *image.Image {
	scope = tg.NewScope(scope)
	contents := op.ReadFile(scope.SubScope("ReadFile"), op.Const(scope.SubScope("filename"), imagePath))
	output := op.DecodePng(scope.SubScope("DecodeJpeg"), contents, op.DecodePngChannels(channels))
	output = op.ExpandDims(scope.SubScope("ExpandDims"), output, op.Const(scope.SubScope("axis"), []int32{0}))
	image := &image.Image{
		Tensor: tg.NewTensor(scope, output)}
	return image.Scale(0, 255)
}

func runModelNormal(base BasePack, model *BaseModel, def_id string, inputImage *tf.Tensor) (order int, confidence float32, err error) {
	order = 0
	err = nil

	tf_model := tg.LoadModel(path.Join("savedData", model.Id, "defs", def_id, "model"), []string{"serve"}, nil)

	results := tf_model.Exec([]tf.Output{
		tf_model.Op("StatefulPartitionedCall", 0),
	}, map[tf.Output]*tf.Tensor{
		tf_model.Op("serving_default_rescaling_input", 0): inputImage,
	})

	var vmax float32 = 0.0
	var predictions = results[0].Value().([][]float32)[0]

	for i, v := range predictions {
		if v > vmax {
			order = i
			vmax = v
		}
	}

	confidence = vmax

	return
}

func runModelExp(base BasePack, model *BaseModel, def_id string, inputImage *tf.Tensor) (order int, confidence float32, err error) {

	err = nil
	order = 0

	base_model := tg.LoadModel(path.Join("savedData", model.Id, "defs", def_id, "base", "model"), []string{"serve"}, nil)

	//results := base_model.Exec([]tf.Output{
	base_results := base_model.Exec([]tf.Output{
		base_model.Op("StatefulPartitionedCall", 0),
	}, map[tf.Output]*tf.Tensor{
		//base_model.Op("serving_default_rescaling_input", 0): inputImage,
		base_model.Op("serving_default_input_1", 0): inputImage,
	})

	type head struct {
		Id          string
		Range_start int
	}

	heads, err := GetDbMultitple[head](base.GetDb(), "exp_model_head where def_id=$1;", def_id)
	if err != nil {
		return
	}

	base.GetLogger().Info("test", "count", len(heads))

	var vmax float32 = 0.0

	for _, element := range heads {
		head_model := tg.LoadModel(path.Join("savedData", model.Id, "defs", def_id, "head", element.Id, "model"), []string{"serve"}, nil)

		results := head_model.Exec([]tf.Output{
			head_model.Op("StatefulPartitionedCall", 0),
		}, map[tf.Output]*tf.Tensor{
			head_model.Op("serving_default_head_input", 0): base_results[0],
		})

		var predictions = results[0].Value().([][]float32)[0]

		for i, v := range predictions {
			base.GetLogger().Debug("predictions", "class", i, "preds", v)
			if v > vmax {
				order = element.Range_start + i
				vmax = v
			}
		}
	}

	// TODO runthe head model
	confidence = vmax

	base.GetLogger().Debug("Got", "heads", len(heads), "order", order, "vmax", vmax)
	return
}

func ClassifyTask(base BasePack, task Task) (err error) {
	defer func() {
		if r := recover(); r != nil {
			base.GetLogger().Error("Task failed due to", "error", r)
			task.UpdateStatusLog(base, TASK_FAILED_RUNNING, "Task failed running")
		}
	}()

	task.UpdateStatusLog(base, TASK_RUNNING, "Runner running task")

	model, err := GetBaseModel(base.GetDb(), *task.ModelId)
	if err != nil {
		task.UpdateStatusLog(base, TASK_FAILED_RUNNING, "Failed to obtain the model")
		return err
	}

	if !model.CanEval() {
		task.UpdateStatusLog(base, TASK_FAILED_RUNNING, "Failed to obtain the model")
		return errors.New("Model not in the right state for evaluation")
	}

	def := JustId{}
	err = GetDBOnce(base.GetDb(), &def, "model_definition where model_id=$1", model.Id)
	if err != nil {
		task.UpdateStatusLog(base, TASK_FAILED_RUNNING, "Failed to obtain the model")
		return
	}

	def_id := def.Id

	// TODO create a database table with tasks
	run_path := path.Join("/tmp", model.Id, "runs")
	os.MkdirAll(run_path, os.ModePerm)

	img_path := path.Join("savedData", model.Id, "tasks", task.Id+"."+model.Format)

	root := tg.NewRoot()

	var tf_img *image.Image = nil

	switch model.Format {
	case "png":
		tf_img = ReadPNG(root, img_path, int64(model.ImageMode))
	case "jpeg":
		tf_img = ReadJPG(root, img_path, int64(model.ImageMode))
	default:
		task.UpdateStatusLog(base, TASK_FAILED_RUNNING, "Failed to obtain the model")
	}

	exec_results := tg.Exec(root, []tf.Output{tf_img.Value()}, nil, &tf.SessionOptions{})
	inputImage, err := tf.NewTensor(exec_results[0].Value())
	if err != nil {
		task.UpdateStatusLog(base, TASK_FAILED_RUNNING, "Failed to run model")
		return
	}

	vi := -1
	var confidence float32 = 0

	if model.ModelType == 2 {
		base.GetLogger().Info("Running model normal", "model", model.Id, "def", def_id)
		vi, confidence, err = runModelExp(base, model, def_id, inputImage)
		if err != nil {
			task.UpdateStatusLog(base, TASK_FAILED_RUNNING, "Failed to run model")
			return
		}
	} else {
		base.GetLogger().Info("Running model normal", "model", model.Id, "def", def_id)
		vi, confidence, err = runModelNormal(base, model, def_id, inputImage)
		if err != nil {
			task.UpdateStatusLog(base, TASK_FAILED_RUNNING, "Failed to run model")
			return
		}
	}

	var GetName struct {
		Name string
		Id   string
	}
	err = GetDBOnce(base.GetDb(), &GetName, "model_classes where model_id=$1 and class_order=$2;", model.Id, vi)
	if err != nil {
		task.UpdateStatusLog(base, TASK_FAILED_RUNNING, "Failed to obtain model results")
		return
	}

	returnValue := struct {
		ClassId    string  `json:"class_id"`
		Class      string  `json:"class"`
		Confidence float32 `json:"confidence"`
	}{
		Class:      GetName.Name,
		ClassId:    GetName.Id,
		Confidence: confidence,
	}

	err = task.SetResult(base, returnValue)
	if err != nil {
		task.UpdateStatusLog(base, TASK_FAILED_RUNNING, "Failed to save model results")
		return
	}

	task.UpdateStatusLog(base, TASK_DONE, "Model ran successfully")

	return
}