fyp/logic/models/data.go

package models

import (
	"archive/zip"
	"bytes"
	"fmt"
	"io"
	"net/http"
	"os"
	"path"
	"reflect"
	"sort"
	"strings"

	model_classes "git.andr3h3nriqu3s.com/andr3/fyp/logic/models/classes"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/models/utils"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/utils"
)

func InsertIfNotPresent(ss []string, s string) []string {
	i := sort.SearchStrings(ss, s)
	if len(ss) > i && ss[i] == s {
		return ss
	}
	ss = append(ss, "")
	copy(ss[i+1:], ss[i:])
	ss[i] = s
	return ss
}

func processZipFile(c *Context, model *BaseModel) {
	reader, err := zip.OpenReader(path.Join("savedData", model.Id, "base_data.zip"))
	if err != nil {
		// TODO add msg to error
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_ZIP_FILE)
		fmt.Printf("Faield to proccess zip file failed to open reader\n")
		fmt.Println(err)
		return
	}
	defer reader.Close()

	training := []string{}
	testing := []string{}

	for _, file := range reader.Reader.File {

		paths := strings.Split(file.Name, "/")

		if paths[1] == "" {
			continue
		}

		if paths[0] != "training" && paths[0] != "testing" {
			fmt.Printf("Invalid file '%s' TODO add msg to response!!!\n", file.Name)
			ModelUpdateStatus(c, model.Id, FAILED_PREPARING_ZIP_FILE)
			return
		}

		if paths[0] != "training" {
			training = InsertIfNotPresent(training, paths[1])
		} else if paths[0] != "testing" {
			testing = InsertIfNotPresent(testing, paths[1])
		}
	}

	if !reflect.DeepEqual(testing, training) {
		fmt.Printf("testing and training are diferent\n")
		fmt.Println(testing)
		fmt.Println(training)
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_ZIP_FILE)
		return
	}

	base_path := path.Join("savedData", model.Id, "data")
	if err = os.MkdirAll(base_path, os.ModePerm); err != nil {
		fmt.Printf("Failed to create base_path dir\n")
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_ZIP_FILE)
		return
	}

	ids := map[string]string{}

	for i, name := range training {
		id, err := model_classes.CreateClass(c.Db, model.Id, i, name)
		if err != nil {
			fmt.Printf("Failed to create class '%s' on db\n", name)
			ModelUpdateStatus(c, id, FAILED_PREPARING_ZIP_FILE)
			return
		}
		ids[name] = id
	}

	for _, file := range reader.Reader.File {
		if file.Name[len(file.Name)-1] == '/' {
			continue
		}

		data, err := reader.Open(file.Name)
		if err != nil {
			fmt.Printf("Could not open file in zip %s\n", file.Name)
			ModelUpdateStatus(c, model.Id, FAILED_PREPARING_ZIP_FILE)
			return
		}
		defer data.Close()
		file_data, err := io.ReadAll(data)
		if err != nil {
			fmt.Printf("Could not read file file in zip %s\n", file.Name)
			ModelUpdateStatus(c, model.Id, FAILED_PREPARING_ZIP_FILE)
			return
		}

		// TODO check if the file is a valid photo that matched the defined photo on the database

		parts := strings.Split(file.Name, "/")

		mode := model_classes.DATA_POINT_MODE_TRAINING
		if parts[0] == "testing" {
			mode = model_classes.DATA_POINT_MODE_TESTING
		}

		data_point_id, err := model_classes.AddDataPoint(c.Db, ids[parts[1]], "id://", mode)
		if err != nil {
			fmt.Printf("Failed to add data point for %s\n", model.Id)
			fmt.Println(err)
			ModelUpdateStatus(c, model.Id, FAILED_PREPARING_ZIP_FILE)
			return
		}

		file_path := path.Join(base_path, data_point_id+"."+model.Format)
		f, err := os.Create(file_path)
		if err != nil {
			fmt.Printf("Could not create file %s\n", file_path)
			ModelUpdateStatus(c, model.Id, FAILED_PREPARING_ZIP_FILE)
			return
		}
		defer f.Close()
		f.Write(file_data)

		if !testImgForModel(c, model, file_path) {
			c.Logger.Errorf("Image did not have valid format for model %s (in zip: %s)!", file_path, file.Name)
			c.Logger.Warn("Not failling updating data point to status -1")
			message := "Image did not have valid format for the model"
			if err = model_classes.UpdateDataPointStatus(c.Db, data_point_id, -1, &message); err != nil {
				c.Logger.Error("Failed to update data point status")
				ModelUpdateStatus(c, model.Id, FAILED_PREPARING_ZIP_FILE)
			}
		}
	}

	fmt.Printf("Added data to model '%s'!\n", model.Id)
	ModelUpdateStatus(c, model.Id, CONFIRM_PRE_TRAINING)
}

func processZipFileExpand(c *Context, model *BaseModel) {
	var err error

	failed := func(msg string) {
		c.Logger.Error(msg, "err", err)
		ModelUpdateStatus(c, model.Id, READY_FAILED)
	}

	reader, err := zip.OpenReader(path.Join("savedData", model.Id, "expand_data.zip"))
	if err != nil {
		failed("Faield to proccess zip file failed to open reader\n")
		return
	}
	defer reader.Close()

	training := []string{}
	testing := []string{}

	for _, file := range reader.Reader.File {

		paths := strings.Split(file.Name, "/")

		if paths[1] == "" {
			continue
		}

		if paths[0] != "training" && paths[0] != "testing" {
			failed(fmt.Sprintf("Invalid file '%s' TODO add msg to response!!!", file.Name))
			return
		}

		if paths[0] != "training" {
			training = InsertIfNotPresent(training, paths[1])
		} else if paths[0] != "testing" {
			testing = InsertIfNotPresent(testing, paths[1])
		}
	}

	if !reflect.DeepEqual(testing, training) {
		failed("testing and training are diferent")
		return
	}

	base_path := path.Join("savedData", model.Id, "data")
	if err = os.MkdirAll(base_path, os.ModePerm); err != nil {
		failed("Failed to create base_path dir")
		return
	}

	ids := map[string]string{}

	for i, name := range training {
		id, err := model_classes.CreateClass(c.Db, model.Id, i, name)
		if err != nil {
			failed(fmt.Sprintf("Failed to create class '%s' on db\n", name))
			return
		}
		ids[name] = id
	}

	for _, file := range reader.Reader.File {
		if file.Name[len(file.Name)-1] == '/' {
			continue
		}

		data, err := reader.Open(file.Name)
		if err != nil {
			failed(fmt.Sprintf("Could not open file in zip %s\n", file.Name))
			return
		}
		defer data.Close()
		file_data, err := io.ReadAll(data)
		if err != nil {
			failed(fmt.Sprintf("Could not read file file in zip %s\n", file.Name))
			return
		}

		// TODO check if the file is a valid photo that matched the defined photo on the database

		parts := strings.Split(file.Name, "/")

		mode := model_classes.DATA_POINT_MODE_TRAINING
		if parts[0] == "testing" {
			mode = model_classes.DATA_POINT_MODE_TESTING
		}

		data_point_id, err := model_classes.AddDataPoint(c.Db, ids[parts[1]], "id://", mode)
		if err != nil {
			failed(fmt.Sprintf("Failed to add data point for %s\n", model.Id))
			return
		}

		file_path := path.Join(base_path, data_point_id+"."+model.Format)
		f, err := os.Create(file_path)
		if err != nil {
			failed(fmt.Sprintf("Could not create file %s\n", file_path))
			return
		}
		defer f.Close()
		f.Write(file_data)

		if !testImgForModel(c, model, file_path) {
			c.Logger.Errorf("Image did not have valid format for model %s (in zip: %s)!", file_path, file.Name)
			c.Logger.Warn("Not failling updating data point to status -1")
			message := "Image did not have valid format for the model"
			if err = model_classes.UpdateDataPointStatus(c.Db, data_point_id, -1, &message); err != nil {
				failed(fmt.Sprintf("Failed to update data point status"))
				return
			}
		}
	}

	c.Logger.Info("Added data to model", "id", model.Id)
	ModelUpdateStatus(c, model.Id, READY)
}

func handleDataUpload(handle *Handle) {
	handle.Post("/models/data/upload", func(c *Context) *Error {
		if !c.CheckAuthLevel(1) {
			return nil
		}

		read_form, err := c.R.MultipartReader()
		if err != nil {
			return c.JsonBadRequest("Please provide a valid form data request!")
		}

		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 c.JsonBadRequest("Please provide a valid form data request!")
			}
			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 c.SendJSONStatus(http.StatusNotFound, "Model not found")
		} else if err != nil {
			return c.Error500(err)
		}

		// TODO mk this path configurable
		dir_path := path.Join("savedData", id)

		f, err := os.Create(path.Join(dir_path, "base_data.zip"))
		if err != nil {
			return c.Error500(err)
		}
		defer f.Close()

		f.Write(file)

		ModelUpdateStatus(c, id, PREPARING_ZIP_FILE)

		go processZipFile(c, model)

		return c.SendJSON(model.Id)
	})

	// ------
	// ------ CLASS DATA UPLOAD
	// ------
	handle.Post("/models/data/class/upload", func(c *Context) *Error {
		if !c.CheckAuthLevel(1) {
			return nil
		}

		read_form, err := c.R.MultipartReader()
		if err != nil {
			return c.JsonBadRequest("Please provide a valid form data request!")
		}

		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 c.JsonBadRequest("Please provide a valid form data request!")
			}
			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()
			}
		}

		c.Logger.Info("Trying to expand model", "id", id)

		model, err := GetBaseModel(handle.Db, id)
		if err == ModelNotFoundError {
			return c.SendJSONStatus(http.StatusNotFound, "Model not found")
		} else if err != nil {
			return c.Error500(err)
		}

		// TODO work in allowing the model to add new in the pre ready moment
		if model.Status != READY {
			return c.JsonBadRequest("Model not in the correct state to add a more classes")
		}

		// TODO mk this path configurable
		dir_path := path.Join("savedData", id)

		f, err := os.Create(path.Join(dir_path, "expand_data.zip"))
		if err != nil {
			return c.Error500(err)
		}
		defer f.Close()

		f.Write(file)

		ModelUpdateStatus(c, id, READY_ALTERATION)

		go processZipFileExpand(c, model)

		return c.SendJSON(model.Id)
	})

	handle.Delete("/models/data/delete-zip-file", func(c *Context) *Error {
		if !c.CheckAuthLevel(1) {
			return nil
		}

		type ModelData struct {
			Id string `json:"id"`
		}

		var dat ModelData

		if err := c.ToJSON(&dat); err != nil {
			return err
		}

		model, err := GetBaseModel(handle.Db, dat.Id)
		if err == ModelNotFoundError {
			return c.SendJSONStatus(http.StatusNotFound, "Model not found")
		} else if err != nil {
			return c.Error500(err)
		}

		delete_path := "base_data.zip"

		if model.Status == READY_FAILED {
			delete_path = "expand_data.zip"
		} else if model.Status != FAILED_PREPARING_ZIP_FILE {
			return c.JsonBadRequest("Model not in the correct status")
		}

		err = os.Remove(path.Join("savedData", model.Id, delete_path))
		if err != nil {
			return c.Error500(err)
		}

		if model.Status != READY_FAILED {
			err = os.RemoveAll(path.Join("savedData", model.Id, "data"))
			if err != nil {
				return c.Error500(err)
			}
		} else {
			c.Logger.Warn("Handle failed to remove the savedData when deleteing the zip file while expanding")
		}

		if model.Status != READY_FAILED {
			_, err = handle.Db.Exec("delete from model_classes where model_id=$1;", model.Id)
			if err != nil {
				return c.Error500(err)
			}
		} else {
			_, err = handle.Db.Exec("delete from model_classes where model_id=$1 and status=$2;", model.Id, MODEL_CLASS_STATUS_TO_TRAIN)
			if err != nil {
				return c.Error500(err)
			}
		}

		if model.Status != READY_FAILED {
			ModelUpdateStatus(c, model.Id, CONFIRM_PRE_TRAINING)
		} else {
			ModelUpdateStatus(c, model.Id, READY)
		}

		return c.SendJSON(model.Id)
	})
}