package models

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"

	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/models/utils"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/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 handleRun(handle *Handle) {
	handle.Post("/models/run", func(w http.ResponseWriter, r *http.Request, c *Context) *Error {
		if !CheckAuthLevel(1, w, r, c) {
			return nil
		}
		if c.Mode == JSON {
			// TODO improve message
			return ErrorCode(nil, 400, nil)
		}

		read_form, err := r.MultipartReader()
		if err != nil {
			// TODO improve message
			return ErrorCode(nil, 400, nil)
		}

		var id string
		var file []byte

		for {
			part, err_part := read_form.NextPart()
			if err_part == io.EOF {
				break
			} else if err_part != nil {
				return &Error{Code: http.StatusBadRequest}
			}
			if part.FormName() == "id" {
				buf := new(bytes.Buffer)
				buf.ReadFrom(part)
				id = buf.String()
			}
			if part.FormName() == "file" {
				buf := new(bytes.Buffer)
				buf.ReadFrom(part)
				file = buf.Bytes()
			}
		}

		model, err := GetBaseModel(handle.Db, id)
		if err == ModelNotFoundError {
			return ErrorCode(nil, http.StatusNotFound, AnyMap{
				"NotFoundMessage": "Model not found",
				"GoBackLink":      "/models",
			})
		} else if err != nil {
			return Error500(err)
		}

		if model.Status != READY {
			// TODO improve this
			return ErrorCode(nil, 400, c.AddMap(nil))
		}

		definitions_rows, err := handle.Db.Query("select id from model_definition where model_id=$1;", model.Id)
		if err != nil {
			return Error500(err)
		}
		defer definitions_rows.Close()

		if !definitions_rows.Next() {
			// TODO improve this
            fmt.Printf("Could not find definition\n")
			return ErrorCode(nil, 400, c.AddMap(nil))
		}

		var def_id string
		if err = definitions_rows.Scan(&def_id); err != nil {
			return Error500(err)
		}

		// 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(run_path, "img." + model.Format)

		img_file, err := os.Create(img_path)
		if err != nil {
			return Error500(err)
		}
		defer img_file.Close()
		img_file.Write(file)

        if !testImgForModel(c, model, img_path) {
            LoadDefineTemplate(w, "/models/edit.html", "run-model-card", c.AddMap(AnyMap{
                "Model": model,
                "NotFound": false,
                "Result": nil,
                "ImageError": true,
            }))
            return nil
        }

		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:
            panic("Not sure what to do with '" + model.Format + "'")
        }

        exec_results := tg.Exec(root, []tf.Output{tf_img.Value()}, nil, &tf.SessionOptions{})
        inputImage, err:= tf.NewTensor(exec_results[0].Value())
        if err != nil {
            return Error500(err)
        }

        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
        vi := 0
        var predictions = results[0].Value().([][]float32)[0]

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

        os.RemoveAll(run_path)

        rows, err := handle.Db.Query("select name from model_classes where model_id=$1 and class_order=$2;", model.Id, vi)
        if err != nil { return Error500(err) }
        if !rows.Next() {
            LoadDefineTemplate(w, "/models/edit.html", "run-model-card", c.AddMap(AnyMap{
                "Model": model,
                "NotFound": true,
                "Result": nil,
            }))
            return nil
        }

        var name string
        if err = rows.Scan(&name); err != nil { return nil }

        LoadDefineTemplate(w, "/models/edit.html", "run-model-card", c.AddMap(AnyMap{
            "Model": model,
            "Result": name,
        }))
		return nil
	})
}