fyp/logic/models/train/train.go

package models_train

import (
	"database/sql"
	"errors"
	"fmt"
	"io"
	"math"
	"net/http"
	"os"
	"os/exec"
	"path"
	"sort"
	"strconv"
	"text/template"

	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"
	"github.com/charmbracelet/log"
)

const EPOCH_PER_RUN = 20
const MAX_EPOCH = 100

func MakeDefenition(db *sql.DB, model_id string, target_accuracy int) (id string, err error) {
	id = ""
	rows, err := db.Query("insert into model_definition (model_id, target_accuracy) values ($1, $2) returning id;", model_id, target_accuracy)
	if err != nil {
		return
	}
	defer rows.Close()
	if !rows.Next() {
		return id, errors.New("Something wrong!")
	}
	err = rows.Scan(&id)
	return
}

func ModelDefinitionUpdateStatus(c *Context, id string, status ModelDefinitionStatus) (err error) {
	_, err = c.Db.Exec("update model_definition set status = $1 where id = $2", status, id)
	return
}

func UpdateStatus(c *Context, table string, id string, status int) (err error) {
	_, err = c.Db.Exec(fmt.Sprintf("update %s set status = $1 where id = $2", table), status, id)
	return
}

func MakeLayer(db *sql.DB, def_id string, layer_order int, layer_type LayerType, shape string) (err error) {
	_, err = db.Exec("insert into model_definition_layer (def_id, layer_order, layer_type, shape) values ($1, $2, $3, $4)", def_id, layer_order, layer_type, shape)
	return
}

func MakeLayerExpandable(db *sql.DB, def_id string, layer_order int, layer_type LayerType, shape string, exp_type int) (err error) {
	_, err = db.Exec("insert into model_definition_layer (def_id, layer_order, layer_type, shape, exp_type) values ($1, $2, $3, $4, $5)", def_id, layer_order, layer_type, shape, exp_type)
	return
}

func generateCvs(c *Context, run_path string, model_id string) (count int, err error) {

	classes, err := c.Db.Query("select count(*) from model_classes where model_id=$1;", model_id)
	if err != nil {
		return
	}
	defer classes.Close()
	if !classes.Next() {
		return
	}
	if err = classes.Scan(&count); err != nil {
		return
	}

	data, err := c.Db.Query("select mdp.id, mc.class_order, mdp.file_path from model_data_point as mdp inner join model_classes as mc on mc.id = mdp.class_id where mc.model_id = $1 and mdp.model_mode=$2;", model_id, model_classes.DATA_POINT_MODE_TRAINING)
	if err != nil {
		return
	}
	defer data.Close()

	f, err := os.Create(path.Join(run_path, "train.csv"))
	if err != nil {
		return
	}
	defer f.Close()
	f.Write([]byte("Id,Index\n"))

	for data.Next() {
		var id string
		var class_order int
		var file_path string
		if err = data.Scan(&id, &class_order, &file_path); err != nil {
			return
		}
		if file_path == "id://" {
			f.Write([]byte(id + "," + strconv.Itoa(class_order) + "\n"))
		} else {
			return count, errors.New("TODO generateCvs to file_path " + file_path)
		}
	}

	return
}

func trainDefinition(c *Context, model *BaseModel, definition_id string, load_prev bool) (accuracy float64, err error) {
	c.Logger.Warn("About to start training definition")
	accuracy = 0
	layers, err := c.Db.Query("select layer_type, shape from model_definition_layer where def_id=$1 order by layer_order asc;", definition_id)
	if err != nil {
		return
	}
	defer layers.Close()

	type layerrow struct {
		LayerType int
		Shape     string
		LayerNum  int
	}

	got := []layerrow{}
	i := 1

	for layers.Next() {
		var row = layerrow{}
		if err = layers.Scan(&row.LayerType, &row.Shape); err != nil {
			return
		}
		row.Shape = shapeToSize(row.Shape)
		row.LayerNum = 1
		got = append(got, row)
		i = i + 1
	}

	// Generate run folder
	run_path := path.Join("/tmp", model.Id, "defs", definition_id)

	err = os.MkdirAll(run_path, os.ModePerm)
	if err != nil {
		return
	}
	defer os.RemoveAll(run_path)

	_, err = generateCvs(c, run_path, model.Id)
	if err != nil {
		return
	}

	// Create python script
	f, err := os.Create(path.Join(run_path, "run.py"))
	if err != nil {
		return
	}
	defer f.Close()

	tmpl, err := template.New("python_model_template.py").ParseFiles("views/py/python_model_template.py")
	if err != nil {
		return
	}

	// Copy result around
	result_path := path.Join("savedData", model.Id, "defs", definition_id)

	if err = tmpl.Execute(f, AnyMap{
		"Layers":           got,
		"Size":             got[0].Shape,
		"DataDir":          path.Join(getDir(), "savedData", model.Id, "data"),
		"RunPath":          run_path,
		"ColorMode":        model.ImageMode,
		"Model":            model,
		"EPOCH_PER_RUN":    EPOCH_PER_RUN,
		"DefId":            definition_id,
		"LoadPrev":         load_prev,
		"LastModelRunPath": path.Join(getDir(), result_path, "model.keras"),
		"SaveModelPath":    path.Join(getDir(), result_path),
	}); err != nil {
		return
	}

	// Run the command
	out, err := exec.Command("bash", "-c", fmt.Sprintf("cd %s && python run.py", run_path)).CombinedOutput()
	if err != nil {
		c.Logger.Debug(string(out))
		return
	}

	c.Logger.Info("Python finished running")

	if err = os.MkdirAll(result_path, os.ModePerm); err != nil {
		return
	}

	accuracy_file, err := os.Open(path.Join(run_path, "accuracy.val"))
	if err != nil {
		return
	}
	defer accuracy_file.Close()

	accuracy_file_bytes, err := io.ReadAll(accuracy_file)
	if err != nil {
		return
	}

	accuracy, err = strconv.ParseFloat(string(accuracy_file_bytes), 64)
	if err != nil {
		return
	}

	c.Logger.Info("Model finished training!", "accuracy", accuracy)
	return
}

func trainDefinitionExp(c *Context, model *BaseModel, definition_id string, load_prev bool) (accuracy float64, err error) {
	accuracy = 0

	c.Logger.Warn("About to start training definition")

	// Get untrained models heads

	// Status = 2 (INIT)
	rows, err := c.Db.Query("select id, range_start, range_end from exp_model_head where def_id=$1 and (status = 2 or status = 3)", definition_id)
	if err != nil {
		return
	}
	defer rows.Close()

	type ExpHead struct {
		id    string
		start int
		end   int
	}

	exp := ExpHead{}

	if rows.Next() {
		if err = rows.Scan(&exp.id, &exp.start, &exp.end); err != nil {
			return
		}
	} else {
		log.Error("Failed to get the exp head of the model")
		err = errors.New("Failed to get the exp head of the model")
		return
	}

	if rows.Next() {
		log.Error("This training function can only train one model at the time")
		err = errors.New("This training function can only train one model at the time")
		return
	}

	UpdateStatus(c, "exp_model_head", exp.id, MODEL_DEFINITION_STATUS_TRAINING)

	layers, err := c.Db.Query("select layer_type, shape, exp_type from model_definition_layer where def_id=$1 order by layer_order asc;", definition_id)
	if err != nil {
		return
	}
	defer layers.Close()

	type layerrow struct {
		LayerType int
		Shape     string
		ExpType   int
		LayerNum  int
	}

	got := []layerrow{}

	remove_top_count := 1

	i := 1

	for layers.Next() {
		var row = layerrow{}
		if err = layers.Scan(&row.LayerType, &row.Shape, &row.ExpType); err != nil {
			return
		}
		row.LayerNum = i
		if row.ExpType == 2 {
			remove_top_count += 1
		}
		row.Shape = shapeToSize(row.Shape)
		got = append(got, row)
		i += 1
	}

	got = append(got, layerrow{
		LayerType: LAYER_DENSE,
		Shape:     fmt.Sprintf("%d", exp.end-exp.start+1),
		ExpType:   2,
		LayerNum:  i,
	})

	// Generate run folder
	run_path := path.Join("/tmp", model.Id, "defs", definition_id)

	err = os.MkdirAll(run_path, os.ModePerm)
	if err != nil {
		return
	}
	defer os.RemoveAll(run_path)

	_, err = generateCvs(c, run_path, model.Id)
	if err != nil {
		return
	}

	// TODO update the run script

	// Create python script
	f, err := os.Create(path.Join(run_path, "run.py"))
	if err != nil {
		return
	}
	defer f.Close()

	tmpl, err := template.New("python_model_template.py").ParseFiles("views/py/python_model_template.py")
	if err != nil {
		return
	}

	// Copy result around
	result_path := path.Join("savedData", model.Id, "defs", definition_id)

	if err = tmpl.Execute(f, AnyMap{
		"Layers":           got,
		"Size":             got[0].Shape,
		"DataDir":          path.Join(getDir(), "savedData", model.Id, "data"),
		"HeadId":           exp.id,
		"RunPath":          run_path,
		"ColorMode":        model.ImageMode,
		"Model":            model,
		"EPOCH_PER_RUN":    EPOCH_PER_RUN,
		"DefId":            definition_id,
		"LoadPrev":         load_prev,
		"LastModelRunPath": path.Join(getDir(), result_path, "model.keras"),
		"SaveModelPath":    path.Join(getDir(), result_path),
		"RemoveTopCount":   remove_top_count,
	}); err != nil {
		return
	}

	// Run the command
	out, err := exec.Command("bash", "-c", fmt.Sprintf("cd %s && python run.py", run_path)).CombinedOutput()
	if err != nil {
		c.Logger.Debug(string(out))
		return
	}

	c.Logger.Info("Python finished running")

	if err = os.MkdirAll(result_path, os.ModePerm); err != nil {
		return
	}

	accuracy_file, err := os.Open(path.Join(run_path, "accuracy.val"))
	if err != nil {
		return
	}
	defer accuracy_file.Close()

	accuracy_file_bytes, err := io.ReadAll(accuracy_file)
	if err != nil {
		return
	}

	accuracy, err = strconv.ParseFloat(string(accuracy_file_bytes), 64)
	if err != nil {
		return
	}

	c.Logger.Info("Model finished training!", "accuracy", accuracy)
	return
}

func remove[T interface{}](lst []T, i int) []T {
	lng := len(lst)
	if i >= lng {
		return []T{}
	}

	if i+1 >= lng {
		return lst[:lng-1]
	}

	if i == 0 {
		return lst[1:]
	}

	return append(lst[:i], lst[i+1:]...)
}

type TrainModelRow struct {
	id              string
	target_accuracy int
	epoch           int
	acuracy         float64
}

type TraingModelRowDefinitions []TrainModelRow

func (nf TraingModelRowDefinitions) Len() int      { return len(nf) }
func (nf TraingModelRowDefinitions) Swap(i, j int) { nf[i], nf[j] = nf[j], nf[i] }
func (nf TraingModelRowDefinitions) Less(i, j int) bool {
	return nf[i].acuracy < nf[j].acuracy
}

type ToRemoveList []int

func (nf ToRemoveList) Len() int      { return len(nf) }
func (nf ToRemoveList) Swap(i, j int) { nf[i], nf[j] = nf[j], nf[i] }
func (nf ToRemoveList) Less(i, j int) bool {
	return nf[i] < nf[j]
}

func trainModel(c *Context, model *BaseModel) {
	definitionsRows, err := c.Db.Query("select id, target_accuracy, epoch from model_definition where status=$1 and model_id=$2", MODEL_DEFINITION_STATUS_INIT, model.Id)
	if err != nil {
		c.Logger.Error("Failed to trainModel!Err:")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}
	defer definitionsRows.Close()

	var definitions TraingModelRowDefinitions = []TrainModelRow{}

	for definitionsRows.Next() {
		var rowv TrainModelRow
		rowv.acuracy = 0
		if err = definitionsRows.Scan(&rowv.id, &rowv.target_accuracy, &rowv.epoch); err != nil {
			c.Logger.Error("Failed to train Model Could not read definition from db!Err:")
			c.Logger.Error(err)
			ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
			return
		}
		definitions = append(definitions, rowv)
	}

	if len(definitions) == 0 {
		c.Logger.Error("No Definitions defined!")
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	firstRound := true
	finished := false

	for {
		var toRemove ToRemoveList = []int{}
		for i, def := range definitions {
			ModelDefinitionUpdateStatus(c, def.id, MODEL_DEFINITION_STATUS_TRAINING)
			accuracy, err := trainDefinition(c, model, def.id, !firstRound)
			if err != nil {
				c.Logger.Error("Failed to train definition!Err:", "err", err)
				ModelDefinitionUpdateStatus(c, def.id, MODEL_DEFINITION_STATUS_FAILED_TRAINING)
				toRemove = append(toRemove, i)
				continue
			}
			def.epoch += EPOCH_PER_RUN
			accuracy = accuracy * 100
			def.acuracy = float64(accuracy)

			definitions[i].epoch += EPOCH_PER_RUN
			definitions[i].acuracy = accuracy

			if accuracy >= float64(def.target_accuracy) {
				c.Logger.Info("Found a definition that reaches target_accuracy!")
				_, err = c.Db.Exec("update model_definition set accuracy=$1, status=$2, epoch=$3 where id=$4", accuracy, MODEL_DEFINITION_STATUS_TRANIED, def.epoch, def.id)
				if err != nil {
					c.Logger.Error("Failed to train definition!Err:\n", "err", err)
					ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
					return
				}

				_, err = c.Db.Exec("update model_definition set status=$1 where id!=$2 and model_id=$3 and status!=$4", MODEL_DEFINITION_STATUS_CANCELD_TRAINING, def.id, model.Id, MODEL_DEFINITION_STATUS_FAILED_TRAINING)
				if err != nil {
					c.Logger.Error("Failed to train definition!Err:\n", "err", err)
					ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
					return
				}

				finished = true
				break
			}

			if def.epoch > MAX_EPOCH {
				fmt.Printf("Failed to train definition! Accuracy less %f < %d\n", accuracy, def.target_accuracy)
				ModelDefinitionUpdateStatus(c, def.id, MODEL_DEFINITION_STATUS_FAILED_TRAINING)
				toRemove = append(toRemove, i)
				continue
			}

			_, err = c.Db.Exec("update model_definition set accuracy=$1, epoch=$2, status=$3 where id=$4", accuracy, def.epoch, MODEL_DEFINITION_STATUS_PAUSED_TRAINING, def.id)
			if err != nil {
				c.Logger.Error("Failed to train definition!Err:\n", "err", err)
				ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
				return
			}
		}

		firstRound = false
		if finished {
			break
		}

		sort.Sort(sort.Reverse(toRemove))

		c.Logger.Info("Round done", "toRemove", toRemove)

		for _, n := range toRemove {
			definitions = remove(definitions, n)
		}

		len_def := len(definitions)

		if len_def == 0 {
			break
		}

		if len_def == 1 {
			continue
		}

		sort.Sort(sort.Reverse(definitions))

		acc := definitions[0].acuracy - 20.0

		c.Logger.Info("Training models, Highest acc", "acc", definitions[0].acuracy, "mod_acc", acc)

		toRemove = []int{}
		for i, def := range definitions {
			if def.acuracy < acc {
				toRemove = append(toRemove, i)
			}
		}

		c.Logger.Info("Removing due to accuracy", "toRemove", toRemove)

		sort.Sort(sort.Reverse(toRemove))
		for _, n := range toRemove {
			c.Logger.Warn("Removing definition not fast enough learning", "n", n)
			ModelDefinitionUpdateStatus(c, definitions[n].id, MODEL_DEFINITION_STATUS_FAILED_TRAINING)
			definitions = remove(definitions, n)
		}
	}

	rows, err := c.Db.Query("select id from model_definition where model_id=$1 and status=$2 order by accuracy desc limit 1;", model.Id, MODEL_DEFINITION_STATUS_TRANIED)
	if err != nil {
		c.Logger.Error("DB: failed to read definition")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}
	defer rows.Close()

	if !rows.Next() {
		// TODO Make the Model status have a message
		c.Logger.Error("All definitions failed to train!")
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	var id string
	if err = rows.Scan(&id); err != nil {
		c.Logger.Error("Failed to read id:")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	if _, err = c.Db.Exec("update model_definition set status=$1 where id=$2;", MODEL_DEFINITION_STATUS_READY, id); err != nil {
		c.Logger.Error("Failed to update model definition")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	to_delete, err := c.Db.Query("select id from model_definition where status != $1 and model_id=$2", MODEL_DEFINITION_STATUS_READY, model.Id)
	if err != nil {
		c.Logger.Error("Failed to select model_definition to delete")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}
	defer to_delete.Close()

	for to_delete.Next() {
		var id string
		if to_delete.Scan(&id); err != nil {
			c.Logger.Error("Failed to scan the id of a model_definition to delete")
			c.Logger.Error(err)
			ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
			return
		}
		os.RemoveAll(path.Join("savedData", model.Id, "defs", id))
	}

	// TODO Check if returning also works here
	if _, err = c.Db.Exec("delete from model_definition where status!=$1 and model_id=$2;", MODEL_DEFINITION_STATUS_READY, model.Id); err != nil {
		c.Logger.Error("Failed to delete model_definition")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	ModelUpdateStatus(c, model.Id, READY)
}

func trainModelExp(c *Context, model *BaseModel) {
	var err error = nil

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

	definitionsRows, err := c.Db.Query("select id, target_accuracy, epoch from model_definition where status=$1 and model_id=$2", MODEL_DEFINITION_STATUS_INIT, model.Id)
	if err != nil {
		failed("Failed to trainModel!")
		return
	}
	defer definitionsRows.Close()

	var definitions TraingModelRowDefinitions = []TrainModelRow{}

	for definitionsRows.Next() {
		var rowv TrainModelRow
		rowv.acuracy = 0
		if err = definitionsRows.Scan(&rowv.id, &rowv.target_accuracy, &rowv.epoch); err != nil {
			failed("Failed to train Model Could not read definition from db!")
			return
		}
		definitions = append(definitions, rowv)
	}

	if len(definitions) == 0 {
		failed("No Definitions defined!")
		return
	}

	firstRound := true
	finished := false

	for {
		var toRemove ToRemoveList = []int{}
		for i, def := range definitions {
			ModelDefinitionUpdateStatus(c, def.id, MODEL_DEFINITION_STATUS_TRAINING)
			accuracy, err := trainDefinitionExp(c, model, def.id, !firstRound)
			if err != nil {
				c.Logger.Error("Failed to train definition!Err:", "err", err)
				ModelDefinitionUpdateStatus(c, def.id, MODEL_DEFINITION_STATUS_FAILED_TRAINING)
				toRemove = append(toRemove, i)
				continue
			}
			def.epoch += EPOCH_PER_RUN
			accuracy = accuracy * 100
			def.acuracy = float64(accuracy)

			definitions[i].epoch += EPOCH_PER_RUN
			definitions[i].acuracy = accuracy

			if accuracy >= float64(def.target_accuracy) {
				c.Logger.Info("Found a definition that reaches target_accuracy!")
				_, err = c.Db.Exec("update model_definition set accuracy=$1, status=$2, epoch=$3 where id=$4", accuracy, MODEL_DEFINITION_STATUS_TRANIED, def.epoch, def.id)
				if err != nil {
					failed("Failed to train definition!")
					return
				}

				_, err = c.Db.Exec("update model_definition set status=$1 where id!=$2 and model_id=$3 and status!=$4", MODEL_DEFINITION_STATUS_CANCELD_TRAINING, def.id, model.Id, MODEL_DEFINITION_STATUS_FAILED_TRAINING)
				if err != nil {
					failed("Failed to train definition!")
					return
				}

				finished = true
				break
			}

			if def.epoch > MAX_EPOCH {
				fmt.Printf("Failed to train definition! Accuracy less %f < %d\n", accuracy, def.target_accuracy)
				ModelDefinitionUpdateStatus(c, def.id, MODEL_DEFINITION_STATUS_FAILED_TRAINING)
				toRemove = append(toRemove, i)
				continue
			}

			_, err = c.Db.Exec("update model_definition set accuracy=$1, epoch=$2, status=$3 where id=$4", accuracy, def.epoch, MODEL_DEFINITION_STATUS_PAUSED_TRAINING, def.id)
			if err != nil {
				failed("Failed to train definition!")
				return
			}
		}

		firstRound = false
		if finished {
			break
		}

		sort.Sort(sort.Reverse(toRemove))

		c.Logger.Info("Round done", "toRemove", toRemove)

		for _, n := range toRemove {
			definitions = remove(definitions, n)
		}

		len_def := len(definitions)

		if len_def == 0 {
			break
		}

		if len_def == 1 {
			continue
		}

		sort.Sort(sort.Reverse(definitions))

		acc := definitions[0].acuracy - 20.0

		c.Logger.Info("Training models, Highest acc", "acc", definitions[0].acuracy, "mod_acc", acc)

		toRemove = []int{}
		for i, def := range definitions {
			if def.acuracy < acc {
				toRemove = append(toRemove, i)
			}
		}

		c.Logger.Info("Removing due to accuracy", "toRemove", toRemove)

		sort.Sort(sort.Reverse(toRemove))
		for _, n := range toRemove {
			c.Logger.Warn("Removing definition not fast enough learning", "n", n)
			ModelDefinitionUpdateStatus(c, definitions[n].id, MODEL_DEFINITION_STATUS_FAILED_TRAINING)
			definitions = remove(definitions, n)
		}
	}

	rows, err := c.Db.Query("select id from model_definition where model_id=$1 and status=$2 order by accuracy desc limit 1;", model.Id, MODEL_DEFINITION_STATUS_TRANIED)
	if err != nil {
		failed("DB: failed to read definition")
		return
	}
	defer rows.Close()

	if !rows.Next() {
		failed("All definitions failed to train!")
		return
	}

	var id string
	if err = rows.Scan(&id); err != nil {
		failed("Failed to read id")
		return
	}

	if _, err = c.Db.Exec("update model_definition set status=$1 where id=$2;", MODEL_DEFINITION_STATUS_READY, id); err != nil {
		failed("Failed to update model definition")
		return
	}

	to_delete, err := c.Db.Query("select id from model_definition where status != $1 and model_id=$2", MODEL_DEFINITION_STATUS_READY, model.Id)
	if err != nil {
		failed("Failed to select model_definition to delete")
		return
	}
	defer to_delete.Close()

	for to_delete.Next() {
		var id string
		if to_delete.Scan(&id); err != nil {
			failed("Failed to scan the id of a model_definition to delete")
			return
		}
		os.RemoveAll(path.Join("savedData", model.Id, "defs", id))
	}

	// TODO Check if returning also works here
	if _, err = c.Db.Exec("delete from model_definition where status!=$1 and model_id=$2;", MODEL_DEFINITION_STATUS_READY, model.Id); err != nil {
		failed("Failed to delete model_definition")
		return
	}

	if err = splitModel(c, model); err != nil {
		failed("Failed to split the model")
		return
	}

	// There should only be one def availabale
	def := JustId{}

	if err = GetDBOnce(c, &def, "model_definition where model_id=$1", model.Id); err != nil {
		return
	}

	// Remove the  base model
	c.Logger.Warn("Removing base model for", "model", model.Id, "def", def.Id)
	os.RemoveAll(path.Join("savedData", model.Id, "defs", def.Id, "model"))
	os.RemoveAll(path.Join("savedData", model.Id, "defs", def.Id, "model.keras"))

	ModelUpdateStatus(c, model.Id, READY)
}

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

	def := JustId{}

	if err = GetDBOnce(c, &def, "model_definition where model_id=$1", model.Id); err != nil {
		return
	}

	head := JustId{}

	if err = GetDBOnce(c, &head, "exp_model_head where def_id=$1", def.Id); err != nil {
		return
	}

	// Generate run folder
	run_path := path.Join("/tmp", model.Id, "defs", def.Id)

	err = os.MkdirAll(run_path, os.ModePerm)
	if err != nil {
		return
	}
	// TODO reneable it
	// defer os.RemoveAll(run_path)

	// Create python script
	f, err := os.Create(path.Join(run_path, "run.py"))
	if err != nil {
		return
	}
	defer f.Close()

	tmpl, err := template.New("python_split_model_template.py").ParseFiles("views/py/python_split_model_template.py")
	if err != nil {
		return
	}

	// Copy result around
	result_path := path.Join(getDir(), "savedData", model.Id, "defs", def.Id)

	// TODO maybe move this to a select count(*)
	// Get only fixed lawers
	layers, err := c.Db.Query("select exp_type from model_definition_layer where def_id=$1 and exp_type=$2 order by layer_order asc;", def.Id, 1)
	if err != nil {
		return
	}
	defer layers.Close()

	type layerrow struct {
		ExpType int
	}

	count := -1

	for layers.Next() {
		count += 1
	}

	if count == -1 {
		err = errors.New("Can not get layers")
		return
	}

	log.Warn("Spliting model", "def", def.Id, "head", head.Id, "count", count)

	basePath := path.Join(result_path, "base")
	headPath := path.Join(result_path, "head", head.Id)

	if err = os.MkdirAll(basePath, os.ModePerm); err != nil {
		return
	}

	if err = os.MkdirAll(headPath, os.ModePerm); err != nil {
		return
	}

	if err = tmpl.Execute(f, AnyMap{
		"SplitLen":      count,
		"ModelPath":     path.Join(result_path, "model.keras"),
		"BaseModelPath": basePath,
		"HeadModelPath": headPath,
	}); err != nil {
		return
	}

	out, err := exec.Command("bash", "-c", fmt.Sprintf("cd %s && python run.py", run_path)).CombinedOutput()
	if err != nil {
		c.Logger.Debug(string(out))
		return
	}

	c.Logger.Info("Python finished running")

	return
}

func removeFailedDataPoints(c *Context, model *BaseModel) (err error) {
	rows, err := c.Db.Query("select mdp.id from model_data_point as mdp join model_classes as mc on mc.id=mdp.class_id where mc.model_id=$1 and mdp.status=-1;", model.Id)
	if err != nil {
		return
	}
	defer rows.Close()

	base_path := path.Join("savedData", model.Id, "data")

	for rows.Next() {
		var dataPointId string
		err = rows.Scan(&dataPointId)
		if err != nil {
			return
		}

		p := path.Join(base_path, dataPointId+"."+model.Format)

		c.Logger.Warn("Removing image", "path", p)

		err = os.RemoveAll(p)
		if err != nil {
			return
		}
	}

	_, err = c.Db.Exec("delete from model_data_point as mdp using model_classes as mc where mdp.class_id = mc.id and mc.model_id=$1 and mdp.status=-1;", model.Id)
	return
}

// This generates a definition
func generateDefinition(c *Context, model *BaseModel, target_accuracy int, number_of_classes int, complexity int) *Error {
	var err error = nil
	failed := func() *Error {
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
		// TODO improve this response
		return c.Error500(err)
	}

	def_id, err := MakeDefenition(c.Db, model.Id, target_accuracy)
	if err != nil {
		return failed()
	}

	order := 1

	// Note the shape of the first layer defines the import size
	if complexity == 2 {
		// Note the shape for now is no used
		width := int(math.Pow(2, math.Floor(math.Log(float64(model.Width))/math.Log(2.0))))
		height := int(math.Pow(2, math.Floor(math.Log(float64(model.Height))/math.Log(2.0))))
		c.Logger.Warn("Complexity 2 creating model with smaller size", "width", width, "height", height)
		err = MakeLayer(c.Db, def_id, order, LAYER_INPUT, fmt.Sprintf("%d,%d,1", width, height))
		if err != nil {
			return failed()
		}
		order++
	} else {
		err = MakeLayer(c.Db, def_id, order, LAYER_INPUT, fmt.Sprintf("%d,%d,1", model.Width, model.Height))
		if err != nil {
			return failed()
		}
		order++
	}

	if complexity == 0 {

		err = MakeLayer(c.Db, def_id, order, LAYER_FLATTEN, "")
		if err != nil {
			return failed()
		}
		order++

		loop := int(math.Log2(float64(number_of_classes)))
		for i := 0; i < loop; i++ {
			err = MakeLayer(c.Db, def_id, order, LAYER_DENSE, fmt.Sprintf("%d,1", number_of_classes*(loop-i)))
			order++
			if err != nil {
				ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
				// TODO improve this response
				return c.Error500(err)
			}
		}

	} else if complexity == 1 || complexity == 2 {

		loop := int((math.Log(float64(model.Width)) / math.Log(float64(10))))
		if loop == 0 {
			loop = 1
		}
		for i := 0; i < loop; i++ {
			err = MakeLayer(c.Db, def_id, order, LAYER_SIMPLE_BLOCK, "")
			order++
			if err != nil {
				return failed()
			}
		}

		err = MakeLayer(c.Db, def_id, order, LAYER_FLATTEN, "")
		if err != nil {
			return failed()
		}
		order++

		loop = int((math.Log(float64(number_of_classes)) / math.Log(float64(10))) / 2)
		if loop == 0 {
			loop = 1
		}
		for i := 0; i < loop; i++ {
			err = MakeLayer(c.Db, def_id, order, LAYER_DENSE, fmt.Sprintf("%d,1", number_of_classes*(loop-i)))
			order++
			if err != nil {
				return failed()
			}
		}
	} else {
		c.Logger.Error("Unkown complexity", "complexity", complexity)
		return failed()
	}

	err = ModelDefinitionUpdateStatus(c, def_id, MODEL_DEFINITION_STATUS_INIT)
	if err != nil {
		return failed()
	}

	return nil
}

func generateDefinitions(c *Context, model *BaseModel, target_accuracy int, number_of_models int) *Error {
	cls, err := model_classes.ListClasses(c.Db, model.Id)
	if err != nil {
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
		// TODO improve this response
		return c.Error500(err)
	}

	err = removeFailedDataPoints(c, model)
	if err != nil {
		return c.Error500(err)
	}

	cls_len := len(cls)

	if number_of_models == 1 {
		if model.Width < 100 && model.Height < 100 && cls_len < 30 {
			generateDefinition(c, model, target_accuracy, cls_len, 0)
		} else if model.Width > 100 && model.Height > 100 {
			generateDefinition(c, model, target_accuracy, cls_len, 2)
		} else {
			generateDefinition(c, model, target_accuracy, cls_len, 1)
		}
	} else if number_of_models == 3 {
		for i := 0; i < number_of_models; i++ {
			generateDefinition(c, model, target_accuracy, cls_len, i)
		}
	} else {
		// TODO handle incrisea the complexity
		for i := 0; i < number_of_models; i++ {
			generateDefinition(c, model, target_accuracy, cls_len, 0)
		}
	}

	return nil
}

func CreateExpModelHead(c *Context, def_id string, range_start int, range_end int, status ModelDefinitionStatus) (id string, err error) {
	rows, err := c.Db.Query("insert into exp_model_head (def_id, range_start, range_end, status) values ($1, $2, $3, $4) returning id", def_id, range_start, range_end, status)

	if err != nil {
		return
	}
	defer rows.Close()

	if !rows.Next() {
		c.Logger.Error("Could not get status of model definition")
		err = errors.New("Could not get status of model definition")
		return
	}

	err = rows.Scan(&id)
	if err != nil {
		return
	}

	return
}

func ExpModelHeadUpdateStatus(db *sql.DB, id string, status ModelDefinitionStatus) (err error) {
	_, err = db.Exec("update model_definition set status = $1 where id = $2", status, id)
	return
}

// This generates a definition
func generateExpandableDefinition(c *Context, model *BaseModel, target_accuracy int, number_of_classes int, complexity int) *Error {
	c.Logger.Info("Generating expandable new definition for model", "id", model.Id, "complexity", complexity)

	var err error = nil
	failed := func() *Error {
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
		// TODO improve this response
		return c.Error500(err)
	}

	if complexity == 0 {
		return failed()
	}

	def_id, err := MakeDefenition(c.Db, model.Id, target_accuracy)
	if err != nil {
		return failed()
	}

	order := 1

	width := model.Width
	height := model.Height

	// Note the shape of the first layer defines the import size
	if complexity == 2 {
		// Note the shape for now is no used
		width := int(math.Pow(2, math.Floor(math.Log(float64(model.Width))/math.Log(2.0))))
		height := int(math.Pow(2, math.Floor(math.Log(float64(model.Height))/math.Log(2.0))))
		c.Logger.Warn("Complexity 2 creating model with smaller size", "width", width, "height", height)

	}

	err = MakeLayerExpandable(c.Db, def_id, order, LAYER_INPUT, fmt.Sprintf("%d,%d,1", width, height), 1)

	order++

	// handle the errors inside the pervious if block
	if err != nil {
		return failed()
	}

	// Create the blocks
	loop := int((math.Log(float64(model.Width)) / math.Log(float64(10))))

	if model.Width < 50 && model.Height < 50 {
		loop = 0
	}

	log.Info("Size of the simple block", "loop", loop)

	//loop = max(loop, 3)

	for i := 0; i < loop; i++ {
		err = MakeLayerExpandable(c.Db, def_id, order, LAYER_SIMPLE_BLOCK, "", 1)
		order++
		if err != nil {
			return failed()
		}
	}

	// Flatten the blocks into dense
	err = MakeLayerExpandable(c.Db, def_id, order, LAYER_FLATTEN, "", 1)
	if err != nil {
		return failed()
	}
	order++

	// Flatten the blocks into dense
	err = MakeLayerExpandable(c.Db, def_id, order, LAYER_DENSE, fmt.Sprintf("%d,1", number_of_classes*2), 1)
	if err != nil {
		return failed()
	}
	order++

	loop = int((math.Log(float64(number_of_classes)) / math.Log(float64(10))) / 2)

	log.Info("Size of the dense layers", "loop", loop)

	// loop = max(loop, 3)

	for i := 0; i < loop; i++ {
		err = MakeLayer(c.Db, def_id, order, LAYER_DENSE, fmt.Sprintf("%d,1", number_of_classes*(loop-i)))
		order++
		if err != nil {
			return failed()
		}
	}

	_, err = CreateExpModelHead(c, def_id, 0, number_of_classes-1, MODEL_DEFINITION_STATUS_INIT)
	if err != nil {
		return failed()
	}

	err = ModelDefinitionUpdateStatus(c, def_id, MODEL_DEFINITION_STATUS_INIT)
	if err != nil {
		return failed()
	}

	return nil
}

// TODO make this json friendy
func generateExpandableDefinitions(c *Context, model *BaseModel, target_accuracy int, number_of_models int) *Error {
	cls, err := model_classes.ListClasses(c.Db, model.Id)
	if err != nil {
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
		// TODO improve this response
		return c.Error500(err)
	}

	err = removeFailedDataPoints(c, model)
	if err != nil {
		return c.Error500(err)
	}

	cls_len := len(cls)

	if number_of_models == 1 {
		if model.Width > 100 && model.Height > 100 {
			generateExpandableDefinition(c, model, target_accuracy, cls_len, 2)
		} else {
			generateExpandableDefinition(c, model, target_accuracy, cls_len, 1)
		}
	} else if number_of_models == 3 {
		for i := 0; i < number_of_models; i++ {
			generateExpandableDefinition(c, model, target_accuracy, cls_len, i)
		}
	} else {
		// TODO handle incrisea the complexity
		for i := 0; i < number_of_models; i++ {
			generateExpandableDefinition(c, model, target_accuracy, cls_len, 1)
		}
	}

	return nil
}

func handle_models_train_json(w http.ResponseWriter, r *http.Request, c *Context) *Error {
	var dat struct {
		Id             string `json:"id"`
		ModelType      string `json:"model_type"`
		NumberOfModels int    `json:"number_of_models"`
		Accuracy       int    `json:"accuracy"`
	}

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

	if dat.Id == "" {
		return c.JsonBadRequest("Please provide a id")
	}

	modelTypeId := 1
	if dat.ModelType == "expandable" {
		modelTypeId = 2
	} else if dat.ModelType != "simple" {
		return c.JsonBadRequest("Invalid model type!")
	}

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

	if model.Status != CONFIRM_PRE_TRAINING {
		return c.JsonBadRequest("Model in invalid status for training")
	}

	if modelTypeId == 2 {
		full_error := generateExpandableDefinitions(c, model, dat.Accuracy, dat.NumberOfModels)
		if full_error != nil {
			return full_error
		}
	} else {
		full_error := generateDefinitions(c, model, dat.Accuracy, dat.NumberOfModels)
		if full_error != nil {
			return full_error
		}
	}

	if modelTypeId == 2 {
		go trainModelExp(c, model)
	} else {
		go trainModel(c, model)
	}

	_, err = c.Db.Exec("update models set status = $1, model_type = $2 where id = $3", TRAINING, modelTypeId, model.Id)
	if err != nil {
		fmt.Println("Failed to update model status")
		fmt.Println(err)
		// TODO improve this response
		return Error500(err)
	}

    return c.SendJSON(model.Id)
}

func handleTrain(handle *Handle) {
	handle.Post("/models/train", func(w http.ResponseWriter, r *http.Request, c *Context) *Error {
		if !CheckAuthLevel(1, w, r, c) {
			return nil
		}

		if c.Mode == JSON {
			return handle_models_train_json(w, r, c)
		}

		r.ParseForm()
		f := r.Form

		number_of_models := 0
		accuracy := 0

		if !CheckId(f, "id") || CheckEmpty(f, "model_type") || !CheckNumber(f, "number_of_models", &number_of_models) || !CheckNumber(f, "accuracy", &accuracy) {
			// TODO improve this response
			return ErrorCode(nil, 400, c.AddMap(nil))
		}

		id := f.Get("id")
		model_type_id := 1
		model_type_form := f.Get("model_type")

		if model_type_form == "expandable" {
			model_type_id = 2
		} else if model_type_form != "simple" {
			return c.Error400(nil, "Invalid model type!", w, "/models/edit.html", "train-model-card", AnyMap{
				"HasData":      true,
				"ErrorMessage": "Invalid model type!",
			})
		}

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

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

		if model_type_id == 2 {
			full_error := generateExpandableDefinitions(c, model, accuracy, number_of_models)
			if full_error != nil {
				return full_error
			}
		} else {
			full_error := generateDefinitions(c, model, accuracy, number_of_models)
			if full_error != nil {
				return full_error
			}
		}

		if model_type_id == 2 {
			go trainModelExp(c, model)
		} else {
			go trainModel(c, model)
		}

		_, err = c.Db.Exec("update models set status = $1, model_type = $2 where id = $3", TRAINING, model_type_id, model.Id)
		if err != nil {
			fmt.Println("Failed to update model status")
			fmt.Println(err)
			// TODO improve this response
			return Error500(err)
		}
		Redirect("/models/edit?id="+model.Id, c.Mode, w, r)
		return nil
	})

	handle.Get("/model/epoch/update", func(w http.ResponseWriter, r *http.Request, c *Context) *Error {
		// TODO check auth level
		if c.Mode != NORMAL {
			// This should only handle normal requests
			c.Logger.Warn("This function only works with normal")
			return c.UnsafeErrorCode(nil, 400, nil)
		}

		f := r.URL.Query()

		accuracy := 0.0

		if !CheckId(f, "model_id") || !CheckId(f, "definition") || CheckEmpty(f, "epoch") || !CheckFloat64(f, "accuracy", &accuracy) {
			c.Logger.Warn("Invalid: model_id or definition or epoch or accuracy")
			return c.UnsafeErrorCode(nil, 400, nil)
		}

		accuracy = accuracy * 100

		model_id := f.Get("model_id")
		def_id := f.Get("definition")
		epoch, err := strconv.Atoi(f.Get("epoch"))
		if err != nil {
			c.Logger.Warn("Epoch is not a number")
			// No need to improve message because this function is only called internaly
			return c.UnsafeErrorCode(nil, 400, nil)
		}

		rows, err := c.Db.Query("select md.status from model_definition as md where md.model_id=$1 and md.id=$2", model_id, def_id)
		if err != nil {
			return c.Error500(err)
		}
		defer rows.Close()

		if !rows.Next() {
			c.Logger.Error("Could not get status of model definition")
			return c.Error500(nil)
		}

		var status int
		err = rows.Scan(&status)
		if err != nil {
			return c.Error500(err)
		}

		if status != 3 {
			c.Logger.Warn("Definition not on status 3(training)", "status", status)
			// No need to improve message because this function is only called internaly
			return c.UnsafeErrorCode(nil, 400, nil)
		}

		c.Logger.Info("Updated model_definition!", "model", model_id, "progress", epoch, "accuracy", accuracy)

		_, err = c.Db.Exec("update model_definition set epoch_progress=$1, accuracy=$2 where id=$3", epoch, accuracy, def_id)
		if err != nil {
			return c.Error500(err)
		}
		return nil
	})

	handle.Get("/model/head/epoch/update", func(w http.ResponseWriter, r *http.Request, c *Context) *Error {
		// TODO check auth level
		if c.Mode != NORMAL {
			// This should only handle normal requests
			c.Logger.Warn("This function only works with normal")
			return c.UnsafeErrorCode(nil, 400, nil)
		}

		f := r.URL.Query()

		accuracy := 0.0

		if !CheckId(f, "head_id") || CheckEmpty(f, "epoch") || !CheckFloat64(f, "accuracy", &accuracy) {
			c.Logger.Warn("Invalid: model_id or head_id or epoch or accuracy")
			return c.UnsafeErrorCode(nil, 400, nil)
		}

		accuracy = accuracy * 100

		head_id := f.Get("head_id")
		epoch, err := strconv.Atoi(f.Get("epoch"))
		if err != nil {
			c.Logger.Warn("Epoch is not a number")
			// No need to improve message because this function is only called internaly
			return c.UnsafeErrorCode(nil, 400, nil)
		}

		rows, err := c.Db.Query("select hd.status from exp_model_head as hd where hd.id=$1;", head_id)
		if err != nil {
			return c.Error500(err)
		}
		defer rows.Close()

		if !rows.Next() {
			c.Logger.Error("Could not get status of model head")
			return c.Error500(nil)
		}

		var status int
		err = rows.Scan(&status)
		if err != nil {
			return c.Error500(err)
		}

		if status != 3 {
			c.Logger.Warn("Head not on status 3(training)", "status", status)
			// No need to improve message because this function is only called internaly
			return c.UnsafeErrorCode(nil, 400, nil)
		}

		c.Logger.Info("Updated model_head!", "head", head_id, "progress", epoch, "accuracy", accuracy)

		_, err = c.Db.Exec("update exp_model_head set epoch_progress=$1, accuracy=$2 where id=$3", epoch, accuracy, head_id)
		if err != nil {
			return c.Error500(err)
		}
		return nil
	})
}