package models_train

import (
	"errors"
	"fmt"
	"io"
	"math"
	"os"
	"os/exec"
	"path"
	"runtime/debug"
	"sort"
	"strconv"
	"strings"
	"text/template"

	"git.andr3h3nriqu3s.com/andr3/fyp/logic/db"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/db_types"

	my_torch "git.andr3h3nriqu3s.com/andr3/fyp/logic/models/train/torch"
	modelloader "git.andr3h3nriqu3s.com/andr3/fyp/logic/models/train/torch/modelloader"
	my_nn "git.andr3h3nriqu3s.com/andr3/fyp/logic/models/train/torch/nn"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/tasks/utils"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/utils"

	"git.andr3h3nriqu3s.com/andr3/gotch"
	torch "git.andr3h3nriqu3s.com/andr3/gotch/ts"
	"github.com/charmbracelet/log"
	"github.com/goccy/go-json"
)

const EPOCH_PER_RUN = 20
const MAX_EPOCH = 100

func shapeToSize(shape string) string {
	split := strings.Split(shape, ",")
	return strings.Join(split[:len(split)-1], ",")
}

func getDir() string {
	dir, err := os.Getwd()
	if err != nil {
		panic(err)
	}
	return dir
}

func MakeLayer(db db.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 db.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 setModelClassStatus(c BasePack, status ModelClassStatus, filter string, args ...any) (err error) {
	_, err = c.GetDb().Exec(fmt.Sprintf("update model_classes set status=%d where %s", status, filter), args...)
	return
}

func generateCvsExp(c BasePack, run_path string, model_id string, doPanic bool) (count int, err error) {

	db := c.GetDb()

	var co struct {
		Count int `db:"count(*)"`
	}
	err = GetDBOnce(db, &co, "model_classes where model_id=$1 and status=$2;", model_id, CLASS_STATUS_TRAINING)
	if err != nil {
		return
	}
	count = co.Count

	if count == 0 {
		err = setModelClassStatus(c, CLASS_STATUS_TRAINING, "model_id=$1 and status=$2;", model_id, CLASS_STATUS_TO_TRAIN)
		if err != nil {
			return
		}

		if doPanic {
			return 0, errors.New("No model classes available")
		}

		return generateCvsExp(c, run_path, model_id, true)
	}

	data, err := 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 and mc.status=$3;", model_id, DATA_POINT_MODE_TRAINING, CLASS_STATUS_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 BasePack, m *BaseModel, def *Definition, in_model *my_torch.ContainerModel, classes []*ModelClass) (accuracy float64, model *my_torch.ContainerModel, err error) {
	log := c.GetLogger()
	db := c.GetDb()

	log.Warn("About to start training definition")

	model = in_model
	accuracy = 0

	if model == nil {
		var layers []*Layer
		layers, err = def.GetLayers(db, " order by layer_order asc")
		if err != nil {
			return
		}

		model = my_torch.BuildModel(layers, 0, true)

	}

	// TODO Make the runner provide this
	device := gotch.CudaIfAvailable()
	// device := gotch.CPU

	result_path := path.Join(getDir(), "savedData", m.Id, "defs", def.Id)
	err = os.MkdirAll(result_path, os.ModePerm)
	if err != nil {
		return
	}

	/* opt1, err := my_nn.DefaultAdamConfig().Build(model.Vs, 0.001)
		if err != nil {
			return
		}

	    opt1.Debug() */

	//log.Info("\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n")

	// TODO remove this
	model.To(device)
	defer model.To(gotch.CPU)

	var ds *modelloader.Dataset
	ds, err = modelloader.NewDataset(db, m, classes[0].ClassOrder, classes[len(classes)-1].ClassOrder)
	if err != nil {
		return
	}

	opt, err := my_nn.DefaultAdamConfig().Build(model.Vs, 0.001)
	if err != nil {
		return
	}

	for epoch := 0; epoch < EPOCH_PER_RUN; epoch++ {
		var trainIter *torch.Iter2
		trainIter, err = ds.TrainIter(32)
		if err != nil {
			return
		}
		// trainIter.ToDevice(device)

		log.Info("epoch", "epoch", epoch)

		var trainLoss float64 = 0
		var trainCorrect float64 = 0
		ok := true
		for ok {
			var item torch.Iter2Item
			var loss *torch.Tensor
			item, ok = trainIter.Next()
			if !ok {
				continue
			}

			data := item.Data
			data, err = data.ToDevice(device, gotch.Float, false, true, false)
			if err != nil {
				return
			}

			data, err = data.SetRequiresGrad(true, true)
			if err != nil {
				return
			}
			err = data.RetainGrad(false)
			if err != nil {
				return
			}

			pred := model.ForwardT(data, true)
			pred, err = pred.SetRequiresGrad(true, true)
			if err != nil {
				return
			}

			pred.RetainGrad(false)

			label := item.Label
			label, err = label.ToDevice(device, gotch.Float, false, true, false)
			if err != nil {
				return
			}
			label, err = label.SetRequiresGrad(true, true)
			if err != nil {
				return
			}
			label.RetainGrad(false)

			// Calculate loss
			loss, err = pred.BinaryCrossEntropyWithLogits(label, &torch.Tensor{}, &torch.Tensor{}, 1, false)
			if err != nil {
				return
			}
			loss, err = loss.SetRequiresGrad(true, false)
			if err != nil {
				return
			}

			err = opt.ZeroGrad()
			if err != nil {
				return
			}

			err = loss.Backward()
			if err != nil {
				return
			}

			err = opt.Step()
			if err != nil {
				return
			}

			vars := model.Vs.Variables()

			for k, v := range vars {
				var grad *torch.Tensor
				grad, err = v.Grad(false)
				if err != nil {
					return
				}

				grad, err = grad.Abs(false)
				if err != nil {
					return
				}

				grad, err = grad.Max(false)
				if err != nil {
					return
				}

				log.Info("[grad check]", "k", k, "grad", grad.Float64Values())
			}

			trainLoss = loss.Float64Values()[0]

			// Calculate accuracy
			/*var p_pred, p_labels *torch.Tensor
			p_pred, err = pred.Argmax([]int64{1}, true, false)
			if err != nil {
				return
			}

			p_labels, err = item.Label.Argmax([]int64{1}, true, false)
			if err != nil {
				return
			}

			floats := p_pred.Float64Values()
			floats_labels := p_labels.Float64Values()

			for i := range floats {
				if floats[i] == floats_labels[i] {
					trainCorrect += 1
				}
			} */

			// panic("fornow")
		}

		//v := []float64{}

		log.Info("model training epoch done loss", "loss", trainLoss, "correct", trainCorrect, "out", ds.TrainImagesSize, "accuracy", trainCorrect/float64(ds.TrainImagesSize))

		/*correct := int64(0)
		//torch.NoGrad(func() {
		ok = true
		testIter := ds.TestIter(64)
		for ok {
			var item torch.Iter2Item
			item, ok = testIter.Next()
			if !ok {
				continue
			}

			output := model.Forward(item.Data)

			var pred, labels *torch.Tensor
			pred, err = output.Argmax([]int64{1}, true, false)
			if err != nil {
				return
			}

			labels, err = item.Label.Argmax([]int64{1}, true, false)
			if err != nil {
				return
			}

			floats := pred.Float64Values()
			floats_labels := labels.Float64Values()

			for i := range floats {
				if floats[i] == floats_labels[i] {
					correct += 1
				}
			}
		}

		accuracy = float64(correct) / float64(ds.TestImagesSize)

		log.Info("Eval accuracy", "accuracy", accuracy)

		err = def.UpdateAfterEpoch(db, accuracy*100)
		if err != nil {
			return
		}*/
		//})
	}

	err = my_torch.SaveModel(model, path.Join(result_path, "model.dat"))
	if err != nil {
		return
	}

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

func generateCvsExpandExp(c BasePack, run_path string, model_id string, offset int, doPanic bool) (count_re int, err error) {
	l, db := c.GetLogger(), c.GetDb()

	var co struct {
		Count int `db:"count(*)"`
	}
	err = GetDBOnce(db, &co, "model_classes where model_id=$1 and status=$2;", model_id, CLASS_STATUS_TRAINING)
	if err != nil {
		return
	}
	l.Info("test here", "count", co)
	count_re = co.Count
	count := co.Count

	if count == 0 {
		err = setModelClassStatus(c, CLASS_STATUS_TRAINING, "model_id=$1 and status=$2;", model_id, CLASS_STATUS_TO_TRAIN)
		if err != nil {
			return
		} else if doPanic {
			return 0, errors.New("No model classes available")
		}
		return generateCvsExpandExp(c, run_path, model_id, offset, true)
	}

	data, err := 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 and mc.status=$3;", model_id, DATA_POINT_MODE_TRAINING, CLASS_STATUS_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"))

	count = 0

	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-offset) + "\n"))
		} else {
			return count, errors.New("TODO generateCvs to file_path " + file_path)
		}
		count += 1
	}

	//
	// This is to load some extra data so that the model has more things to train on
	//

	data_other, err := 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 and mc.status=$3 limit $4;", model_id, DATA_POINT_MODE_TRAINING, CLASS_STATUS_TRAINED, count*10)
	if err != nil {
		return
	}
	defer data_other.Close()

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

	return
}

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

	l := c.GetLogger()
	l.Warn("About to retrain model")

	// Get untrained models heads

	type ExpHead struct {
		Id    string
		Start int `db:"range_start"`
		End   int `db:"range_end"`
	}

	// status = 2 (INIT) 3 (TRAINING)
	heads, err := GetDbMultitple[ExpHead](c.GetDb(), "exp_model_head where def_id=$1 and (status = 2 or status = 3)", definition_id)
	if err != nil {
		return
	} else if len(heads) == 0 {
		log.Error("Failed to get the exp head of the model")
		return
	} else if len(heads) != 1 {
		err = errors.New("This training function can only train one model at the time")
		return
	}

	exp := heads[0]

	l.Info("Got exp head", "head", exp)

	if err = UpdateStatus(c.GetDb(), "exp_model_head", exp.Id, DEFINITION_STATUS_TRAINING); err != nil {
		return
	}

	layers, err := c.GetDb().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{}

	i := 1
	var last *layerrow = nil
	got_2 := false

	var first *layerrow = nil

	for layers.Next() {
		var row = layerrow{}
		if err = layers.Scan(&row.LayerType, &row.Shape, &row.ExpType); err != nil {
			return
		}

		// Keep track of the first layer so we can keep the size of the image
		if first == nil {
			first = &row
		}

		row.LayerNum = i
		row.Shape = shapeToSize(row.Shape)
		if row.ExpType == 2 {
			if !got_2 {
				got = append(got, *last)
				got_2 = true
			}
			got = append(got, row)
		}
		last = &row
		i += 1
	}

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

	l.Info("Got layers", "layers", got)

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

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

	classCount, err := generateCvsExpandExp(c, run_path, model.Id, exp.Start, false)
	if err != nil {
		return
	}

	l.Info("Generated cvs", "classCount", classCount)

	// 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()

	l.Info("About to run python!")

	tmpl, err := template.New("python_model_template_expand.py").ParseFiles("views/py/python_model_template_expand.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":             first.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,
		"LoadPrev":         load_prev,
		"BaseModel":        path.Join(getDir(), result_path, "base", "model.keras"),
		"LastModelRunPath": path.Join(getDir(), result_path, "head", exp.Id, "model.keras"),
		"SaveModelPath":    path.Join(getDir(), result_path, "head", exp.Id),
		"Depth":            classCount,
		"StartPoint":       0,
		"Host":             c.GetHost(),
	}); 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 {
		l.Warn("Python failed to run", "err", err, "out", string(out))
		return
	}

	l.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
	}

	os.RemoveAll(run_path)
	l.Info("Model finished training!", "accuracy", accuracy)
	return
}

func trainDefinitionExp(c BasePack, model *BaseModel, definition_id string, load_prev bool) (accuracy float64, err error) {
	accuracy = 0
	l := c.GetLogger()
	db := c.GetDb()

	l.Warn("About to start training definition")

	// Get untrained models heads

	type ExpHead struct {
		Id    string
		Start int `db:"range_start"`
		End   int `db:"range_end"`
	}

	// status = 2 (INIT) 3 (TRAINING)
	heads, err := GetDbMultitple[ExpHead](db, "exp_model_head where def_id=$1 and (status = 2 or status = 3)", definition_id)
	if err != nil {
		return
	} else if len(heads) == 0 {
		log.Error("Failed to get the exp head of the model")
		return
	} else if len(heads) != 1 {
		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
	}

	exp := heads[0]

	if err = UpdateStatus(db, "exp_model_head", exp.Id, DEFINITION_STATUS_TRAINING); err != nil {
		return
	}

	layers, err := 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{}
	i := 1

	for layers.Next() {
		var row = layerrow{}
		if err = layers.Scan(&row.LayerType, &row.Shape, &row.ExpType); err != nil {
			return
		}
		row.LayerNum = i
		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
	}

	classCount, err := generateCvsExp(c, run_path, model.Id, false)
	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,
		"LoadPrev":         load_prev,
		"LastModelRunPath": path.Join(getDir(), result_path, "model.keras"),
		"SaveModelPath":    path.Join(getDir(), result_path),
		"Depth":            classCount,
		"StartPoint":       0,
		"Host":             c.GetHost(),
	}); 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 {
		l.Debug(string(out))
		return
	}

	l.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
	}

	os.RemoveAll(run_path)
	l.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 BasePack, model *BaseModel) (err error) {
	db := c.GetDb()
	log := c.GetLogger()

	fail := func(err error) {
		log.Error("Failed to train Model!", "err", err, "stack", string(debug.Stack()))
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
	}

	defs, err := model.GetDefinitions(db, "and md.status=$2", DEFINITION_STATUS_INIT)
	if err != nil {
		fail(err)
		return
	}

	var definitions SortByAccuracyDefinitions = defs

	if len(definitions) == 0 {
		fail(errors.New("No definitons defined!"))
		return
	}

	finished := false

	models := map[string]*my_torch.ContainerModel{}

	classes, err := model.GetClasses(db, " and status=$2 order by mc.class_order asc", CLASS_STATUS_TO_TRAIN)

	for {
		// Keep track of definitions that did not train fast enough
		var toRemove ToRemoveList = []int{}

		for i, def := range definitions {

			err := def.UpdateStatus(c, DEFINITION_STATUS_TRAINING)
			if err != nil {
				log.Error("Could not make model into training", "err", err)
				def.UpdateStatus(c, DEFINITION_STATUS_FAILED_TRAINING)
				toRemove = append(toRemove, i)
				continue
			}

			accuracy, ml_model, err := trainDefinition(c, model, def, models[def.Id], classes)
			if err != nil {
				log.Error("Failed to train definition!Err:", "err", err)
				def.UpdateStatus(c, DEFINITION_STATUS_FAILED_TRAINING)
				toRemove = append(toRemove, i)
				continue
			}
			models[def.Id] = ml_model

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

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

				finished = true
				break
			}

			if def.Epoch > MAX_EPOCH {
				fmt.Printf("Failed to train definition! Accuracy less %f < %d\n", accuracy, def.TargetAccuracy)
				def.UpdateStatus(c, DEFINITION_STATUS_FAILED_TRAINING)
				toRemove = append(toRemove, i)
				continue
			}

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

		if finished {
			break
		}

		sort.Sort(sort.Reverse(toRemove))

		log.Info("Round done", "toRemove", toRemove)

		for _, n := range toRemove {
			// Clean up unsed models
			models[definitions[n].Id] = nil
			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].Accuracy - 20.0

		log.Info("Training models, Highest acc", "acc", definitions[0].Accuracy, "mod_acc", acc)

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

		log.Info("Removing due to accuracy", "toRemove", toRemove)

		sort.Sort(sort.Reverse(toRemove))
		for _, n := range toRemove {
			log.Warn("Removing definition not fast enough learning", "n", n)
			definitions[n].UpdateStatus(c, DEFINITION_STATUS_FAILED_TRAINING)
			models[definitions[n].Id] = nil
			definitions = remove(definitions, n)
		}
	}

	var def Definition
	err = GetDBOnce(c, &def, "model_definition as md where md.model_id=$1 and md.status=$2 order by md.accuracy desc limit 1;", model.Id, DEFINITION_STATUS_TRANIED)
	if err != nil {
		if err == NotFoundError {
			log.Error("All definitions failed to train!")
		} else {
			log.Error("DB: failed to read definition", "err", err)
		}
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	if err = def.UpdateStatus(c, DEFINITION_STATUS_READY); err != nil {
		log.Error("Failed to update model definition", "err", err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

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

	for to_delete.Next() {
		var id string
		if err = to_delete.Scan(&id); err != nil {
			log.Error("Failed to scan the id of a model_definition to delete", "err", 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 = db.Exec("delete from model_definition where status!=$1 and model_id=$2;", DEFINITION_STATUS_READY, model.Id); err != nil {
		log.Error("Failed to delete model_definition")
		log.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	ModelUpdateStatus(c, model.Id, READY)

	return
}

type TrainModelRowUsable struct {
	Id             string
	TargetAccuracy int `db:"target_accuracy"`
	Epoch          int
	Acuracy        float64 `db:"0"`
}

type TrainModelRowUsables []*TrainModelRowUsable

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

func trainModelExp(c BasePack, model *BaseModel) (err error) {
	l := c.GetLogger()
	db := c.GetDb()

	var definitions TrainModelRowUsables

	definitions, err = GetDbMultitple[TrainModelRowUsable](db, "model_definition where status=$1 and model_id=$2", DEFINITION_STATUS_INIT, model.Id)
	if err != nil {
		l.Error("Failed to get definitions")
		return
	}
	if len(definitions) == 0 {
		l.Error("No Definitions defined!")
		return errors.New("No Definitions found")
	}

	firstRound := true
	finished := false

	for {
		var toRemove ToRemoveList = []int{}
		for i, def := range definitions {
			Definition{Id: def.Id}.UpdateStatus(c, DEFINITION_STATUS_TRAINING)
			accuracy, err := trainDefinitionExp(c, model, def.Id, !firstRound)
			if err != nil {
				l.Error("Failed to train definition!Err:", "err", err)
				Definition{Id: def.Id}.UpdateStatus(c, DEFINITION_STATUS_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.TargetAccuracy) {
				l.Info("Found a definition that reaches target_accuracy!")
				_, err = db.Exec("update model_definition set accuracy=$1, status=$2, epoch=$3 where id=$4", accuracy, DEFINITION_STATUS_TRANIED, def.Epoch, def.Id)
				if err != nil {
					l.Error("Failed to train definition!")
					return err
				}

				_, err = db.Exec("update model_definition set status=$1 where id!=$2 and model_id=$3 and status!=$4", DEFINITION_STATUS_CANCELD_TRAINING, def.Id, model.Id, DEFINITION_STATUS_FAILED_TRAINING)
				if err != nil {
					l.Error("Failed to train definition!")
					return err
				}

				_, err = db.Exec("update exp_model_head set status=$1 where def_id=$2;", MODEL_HEAD_STATUS_READY, def.Id)
				if err != nil {
					l.Error("Failed to train definition!")
					return err
				}

				finished = true
				break
			}

			if def.Epoch > MAX_EPOCH {
				fmt.Printf("Failed to train definition! Accuracy less %f < %d\n", accuracy, def.TargetAccuracy)

				Definition{Id: def.Id}.UpdateStatus(c, DEFINITION_STATUS_FAILED_TRAINING)
				toRemove = append(toRemove, i)
				continue
			}

			_, err = db.Exec("update model_definition set accuracy=$1, epoch=$2, status=$3 where id=$4", accuracy, def.Epoch, DEFINITION_STATUS_PAUSED_TRAINING, def.Id)
			if err != nil {
				l.Error("Failed to train definition!")
				return err
			}
		}

		firstRound = false
		if finished {
			break
		}

		sort.Sort(sort.Reverse(toRemove))

		l.Info("Round done", "toRemove", toRemove)

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

		len_def := len(definitions)

		if len_def == 0 {
			break
		} else if len_def == 1 {
			continue
		}

		sort.Sort(sort.Reverse(definitions))
		acc := definitions[0].Acuracy - 20.0

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

		l.Info("Removing due to accuracy", "toRemove", toRemove)

		sort.Sort(sort.Reverse(toRemove))
		for _, n := range toRemove {
			l.Warn("Removing definition not fast enough learning", "n", n)

			Definition{Id: definitions[n].Id}.UpdateStatus(c, DEFINITION_STATUS_FAILED_TRAINING)
			definitions = remove(definitions, n)
		}
	}

	var dat JustId

	err = GetDBOnce(db, &dat, "model_definition where model_id=$1 and status=$2 order by accuracy desc limit 1;", model.Id, DEFINITION_STATUS_TRANIED)
	if err == NotFoundError {
		// Set the class status to trained
		err = setModelClassStatus(c, CLASS_STATUS_TO_TRAIN, "model_id=$1 and status=$2;", model.Id, CLASS_STATUS_TRAINING)
		if err != nil {
			l.Error("All definitions failed to train! And Failed to set class status")
			return err
		}

		l.Error("All definitions failed to train!")
		return err
	} else if err != nil {
		l.Error("All definitions failed to train!")
		return err
	}

	if _, err = db.Exec("update model_definition set status=$1 where id=$2;", DEFINITION_STATUS_READY, dat.Id); err != nil {
		l.Error("Failed to update model definition")
		return err
	}

	to_delete, err := GetDbMultitple[JustId](db, "model_definition where status!=$1 and model_id=$2", DEFINITION_STATUS_READY, model.Id)
	if err != nil {
		l.Error("Failed to select model_definition to delete")
		return err
	}

	for _, d := range to_delete {
		os.RemoveAll(path.Join("savedData", model.Id, "defs", d.Id))
	}

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

	if err = splitModel(c, model); err != nil {
		err = setModelClassStatus(c, CLASS_STATUS_TO_TRAIN, "model_id=$1 and status=$2;", model.Id, CLASS_STATUS_TRAINING)
		if err != nil {
			l.Error("Failed to split the model! And Failed to set class status")
			return err
		}

		l.Error("Failed to split the model")
		return err
	}

	// Set the class status to trained
	err = setModelClassStatus(c, CLASS_STATUS_TRAINED, "model_id=$1 and status=$2;", model.Id, CLASS_STATUS_TRAINING)
	if err != nil {
		l.Error("Failed to set class status")
		return err
	}

	// There should only be one def availabale
	def := JustId{}
	if err = GetDBOnce(db, &def, "model_definition where model_id=$1", model.Id); err != nil {
		return
	}

	// Remove the  base model
	l.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)
	return
}

func splitModel(c BasePack, model *BaseModel) (err error) {
	db := c.GetDb()
	l := c.GetLogger()

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

	head := JustId{}
	if err = GetDBOnce(db, &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+"-split")

	err = os.MkdirAll(run_path, os.ModePerm)
	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_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 := 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 {
		l.Debug(string(out))
		return
	}

	os.RemoveAll(run_path)
	l.Info("Python finished running")
	return
}

// This generates a definition
func generateDefinition(c BasePack, model *BaseModel, target_accuracy int, number_of_classes int, complexity int) (err error) {
	failed := func() {
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
	}

	db := c.GetDb()
	l := c.GetLogger()

	def, err := MakeDefenition(db, model.Id, target_accuracy)
	if err != nil {
		failed()
		return
	}

	order := 1

	// Note the shape of the first layer defines the import size
	//_, err = def.MakeLayer(db, order, LAYER_INPUT, ShapeToString(model.Width, model.Height, model.ImageMode))
	_, err = def.MakeLayer(db, order, LAYER_INPUT, ShapeToString(3, model.Width, model.Height))
	if err != nil {
		failed()
		return
	}
	order++

	if complexity == 0 {
		_, err = def.MakeLayer(db, order, LAYER_FLATTEN, "")
		if err != nil {
			failed()
			return
		}
		order++

		loop := int(math.Log2(float64(number_of_classes)))
		for i := 0; i < loop; i++ {
			_, err = def.MakeLayer(db, order, LAYER_DENSE, ShapeToString(number_of_classes*(loop-i)))
			order++
			if err != nil {
				ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
				return
			}
		}
	} else if complexity == 1 || complexity == 2 {
		loop := max(1, int((math.Log(float64(model.Width)) / math.Log(float64(10)))))
		for i := 0; i < loop; i++ {
			_, err = def.MakeLayer(db, order, LAYER_SIMPLE_BLOCK, "")
			order++
			if err != nil {
				failed()
				return
			}
		}

		_, err = def.MakeLayer(db, order, LAYER_FLATTEN, "")
		if err != nil {
			failed()
			return
		}
		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 = def.MakeLayer(db, order, LAYER_DENSE, ShapeToString(number_of_classes*(loop-i)))
			order++
			if err != nil {
				failed()
				return
			}
		}
	} else {
		l.Error("Unkown complexity", "complexity", complexity)
		failed()
		return
	}

	return def.UpdateStatus(db, DEFINITION_STATUS_INIT)
}

func generateDefinitions(c BasePack, model *BaseModel, target_accuracy int, number_of_models int) (err error) {
	cls, err := model.GetClasses(c, "")
	if err != nil {
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
		return
	}

	cls_len := len(cls)

	if number_of_models == 1 {
		if model.Width < 100 && model.Height < 100 && cls_len < 30 {
			err = generateDefinition(c, model, target_accuracy, cls_len, 0)
		} else if model.Width > 100 && model.Height > 100 {
			err = generateDefinition(c, model, target_accuracy, cls_len, 2)
		} else {
			err = generateDefinition(c, model, target_accuracy, cls_len, 1)
		}
		if err != nil {
			return
		}
	} else {
		for i := 0; i < number_of_models; i++ {
			err = generateDefinition(c, model, target_accuracy, cls_len, min(i, 2))
			if err != nil {
				return
			}
		}
	}

	return nil
}

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

// This generates a definition
func generateExpandableDefinition(c BasePack, model *BaseModel, target_accuracy int, number_of_classes int, complexity int) (err error) {
	l := c.GetLogger()
	db := c.GetDb()
	l.Info("Generating expandable new definition for model", "id", model.Id, "complexity", complexity)

	failed := func() {
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
	}

	if complexity == 0 {
		failed()
		return
	}

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

	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))))
		l.Warn("Complexity 2 creating model with smaller size", "width", width, "height", height)
	}

	err = MakeLayerExpandable(c.GetDb(), 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 {
		failed()
		return
	}

	// 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(db, def.Id, order, LAYER_SIMPLE_BLOCK, "", 1)
		order++
		if err != nil {
			failed()
			return
		}
	}

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

	// Flatten the blocks into dense
	err = MakeLayerExpandable(db, def.Id, order, LAYER_DENSE, fmt.Sprintf("%d,1", number_of_classes*2), 1)
	if err != nil {
		failed()
		return
	}
	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(db, def.Id, order, LAYER_DENSE, fmt.Sprintf("%d,1", number_of_classes*(loop-i)))
		order++
		if err != nil {
			failed()
			return
		}
	}

	var newHead = struct {
		DefId      string           `db:"def_id"`
		RangeStart int              `db:"range_start"`
		RangeEnd   int              `db:"range_end"`
		Status     DefinitionStatus `db:"status"`
	}{
		def.Id, 0, number_of_classes - 1, DEFINITION_STATUS_INIT,
	}
	_, err = InsertReturnId(c.GetDb(), &newHead, "exp_model_head", "id")
	if err != nil {
		failed()
		return
	}

	err = def.UpdateStatus(c, DEFINITION_STATUS_INIT)
	if err != nil {
		failed()
		return
	}

	return
}

// TODO make this json friendy
func generateExpandableDefinitions(c BasePack, model *BaseModel, target_accuracy int, number_of_models int) (err error) {
	cls, err := model.GetClasses(c, "")
	if err != nil {
		ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
		return
	}

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

	return nil
}

func ResetClasses(c BasePack, model *BaseModel) {
	_, err := c.GetDb().Exec("update model_classes set status=$1 where status=$2 and model_id=$3", CLASS_STATUS_TO_TRAIN, CLASS_STATUS_TRAINING, model.Id)
	if err != nil {
		c.GetLogger().Error("Error while reseting the classes", "error", err)
	}
}

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

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

	var definitions TrainModelRowUsables

	definitions, err = GetDbMultitple[TrainModelRowUsable](c, "model_definition where status=$1 and model_id=$2", DEFINITION_STATUS_READY, model.Id)
	if err != nil {
		failed("Failed to get definitions")
		return
	}
	if len(definitions) != 1 {
		failed("There should only be one definition available!")
		return
	}

	firstRound := true
	def := definitions[0]
	epoch := 0

	for {
		acc, err := trainDefinitionExp(c, model, def.Id, !firstRound)
		if err != nil {
			failed("Failed to train definition!")
			return
		}
		epoch += EPOCH_PER_RUN

		if float64(acc*100) >= float64(def.Acuracy) {
			c.Logger.Info("Found a definition that reaches target_accuracy!")

			_, err = c.Db.Exec("update exp_model_head set status=$1 where def_id=$2 and status=$3;", MODEL_HEAD_STATUS_READY, def.Id, MODEL_HEAD_STATUS_TRAINING)
			if err != nil {
				failed("Failed to train definition!")
				return
			}
			break
		} else if def.Epoch > MAX_EPOCH {
			failed(fmt.Sprintf("Failed to train definition! Accuracy less %f < %d\n", acc*100, def.TargetAccuracy))
			return
		}
	}

	// Set the class status to trained
	err = setModelClassStatus(c, CLASS_STATUS_TRAINED, "model_id=$1 and status=$2;", model.Id, CLASS_STATUS_TRAINING)
	if err != nil {
		failed("Failed to set class status")
		return
	}

	ModelUpdateStatus(c, model.Id, READY)
}

func RunTaskTrain(b BasePack, task Task) (err error) {
	l := b.GetLogger()

	model, err := GetBaseModel(b.GetDb(), *task.ModelId)
	if err != nil {
		task.UpdateStatusLog(b, TASK_FAILED_RUNNING, "Failed to get model information")
		l.Error("Failed to get model information", "err", err)
		return err
	} else if model.Status != TRAINING {
		task.UpdateStatusLog(b, TASK_FAILED_RUNNING, "Model not in the correct status for training")
		return errors.New("Model not in the right status")
	}

	task.UpdateStatusLog(b, TASK_RUNNING, "Training model")

	var dat struct {
		NumberOfModels int
		Accuracy       int
	}

	err = json.Unmarshal([]byte(task.ExtraTaskInfo), &dat)
	if err != nil {
		task.UpdateStatusLog(b, TASK_FAILED_RUNNING, "Failed to get model extra information")
	}

	if model.ModelType == 2 {
		task.UpdateStatusLog(b, TASK_FAILED_RUNNING, "TODO expandable models")
		ModelUpdateStatus(b, model.Id, FAILED_TRAINING)
		panic("todo")

		full_error := generateExpandableDefinitions(b, model, dat.Accuracy, dat.NumberOfModels)
		if full_error != nil {
			l.Error("Failed to generate defintions", "err", full_error)
			task.UpdateStatusLog(b, TASK_FAILED_RUNNING, "Failed generate model")
			return errors.New("Failed to generate definitions")
		}
	} else {
		full_error := generateDefinitions(b, model, dat.Accuracy, dat.NumberOfModels)
		if full_error != nil {
			l.Error("Failed to generate defintions", "err", full_error)
			task.UpdateStatusLog(b, TASK_FAILED_RUNNING, "Failed generate model")
			return errors.New("Failed to generate definitions")
		}
	}

	if model.ModelType == 2 {
		err = trainModelExp(b, model)
	} else {
		err = trainModel(b, model)
	}

	if err != nil {
		l.Error("Failed to train model", "err", err)
		task.UpdateStatusLog(b, TASK_FAILED_RUNNING, "Failed generate model")
		ModelUpdateStatus(b, model.Id, FAILED_TRAINING)
		return
	}

	task.UpdateStatusLog(b, TASK_DONE, "Model finished training")

	return
}

func RunTaskRetrain(b BasePack, task Task) (err error) {
	task.UpdateStatusLog(b, TASK_FAILED_RUNNING, "TODO retrain with torch")
	panic("TODO")

	model, err := GetBaseModel(b.GetDb(), *task.ModelId)
	if err != nil {
		return err
	} else if model.Status != READY_RETRAIN {
		return errors.New("Model in invalid status for re-training")
	}

	l := b.GetLogger()
	db := b.GetDb()

	failed := func() {
		ResetClasses(b, model)
		ModelUpdateStatus(b, model.Id, READY_RETRAIN_FAILED)
		task.UpdateStatusLog(b, TASK_FAILED_RUNNING, "Model failed retraining")
		l.Error("Failed to retrain", "err", err)
	}

	task.UpdateStatusLog(b, TASK_RUNNING, "Model retraining")

	var defData struct {
		Id             string  `db:"md.id"`
		TargetAcuuracy float64 `db:"md.target_accuracy"`
	}

	err = GetDBOnce(db, &defData, "models as m inner join model_definition as md on m.id = md.model_id where m.id=$1;", task.ModelId)
	if err != nil {
		failed()
		return
	}

	var acc float64 = 0
	var epocs = 0
	// TODO make max epochs come from db
	for acc*100 < defData.TargetAcuuracy && epocs < 20 {
		// This is something I have to check
		acc, err = trainDefinitionExpandExp(b, model, defData.Id, epocs > 0)
		if err != nil {
			failed()
			return
		}

		l.Info("Retrained model", "accuracy", acc, "target", defData.TargetAcuuracy)

		epocs += 1
	}

	if acc*100 < defData.TargetAcuuracy {
		l.Error("Model never achived targetd accuracy", "acc", acc*100, "target", defData.TargetAcuuracy)
		failed()
		return
	}

	// TODO check accuracy
	err = UpdateStatus(db, "models", model.Id, READY)
	if err != nil {
		failed()
		return
	}

	l.Info("Model updaded")

	_, err = db.Exec("update model_classes set status=$1 where status=$2 and model_id=$3", CLASS_STATUS_TRAINED, CLASS_STATUS_TRAINING, model.Id)
	if err != nil {
		l.Error("Error while updating the classes", "error", err)
		failed()
		return
	}

	task.UpdateStatusLog(b, TASK_DONE, "Model finished retraining")

	return
}

func handleTrain(handle *Handle) {

	type TrainReq struct {
		Id             string `json:"id" validate:"required"`
		ModelType      string `json:"model_type"`
		NumberOfModels int    `json:"number_of_models"`
		Accuracy       int    `json:"accuracy"`
	}
	PostAuthJson(handle, "/models/train", User_Normal, func(c *Context, dat *TrainReq) *Error {
		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 == NotFoundError {
			return c.JsonBadRequest("Model not found")
		} else if err != nil {
			return c.E500M("Failed to get model information", err)
		} else if model.CanTrain == 0 {
			return c.JsonBadRequest("Model can not be trained!")
		}

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

		_, err = c.Db.Exec("update models set status = $1, model_type = $2 where id = $3", TRAINING, modelTypeId, model.Id)
		if err != nil {
			return c.E500M("Failed to update model_status", err)
		}

		text, err := json.Marshal(struct {
			NumberOfModels int
			Accuracy       int
		}{
			NumberOfModels: dat.NumberOfModels,
			Accuracy:       dat.Accuracy,
		})
		if err != nil {
			return c.E500M("Failed create data", err)
		}

		type CreateNewTask struct {
			UserId        string   `db:"user_id"`
			ModelId       string   `db:"model_id"`
			TaskType      TaskType `db:"task_type"`
			Status        int      `db:"status"`
			ExtraTaskInfo string   `db:"extra_task_info"`
		}

		newTask := CreateNewTask{
			UserId:        c.User.Id,
			ModelId:       model.Id,
			TaskType:      TASK_TYPE_TRAINING,
			Status:        1,
			ExtraTaskInfo: string(text),
		}

		id, err := InsertReturnId(c, &newTask, "tasks", "id")
		if err != nil {
			return c.E500M("Failed to create task", err)
		}

		return c.SendJSON(id)
	})

	PostAuthJson(handle, "/model/train/retrain", User_Normal, func(c *Context, dat *JustId) *Error {
		model, err := GetBaseModel(c.Db, dat.Id)
		if err == NotFoundError {
			return c.JsonBadRequest("Model not found")
		} else if err != nil {
			return c.E500M("Faield to get model", err)
		} else if model.Status != READY && model.Status != READY_RETRAIN_FAILED && model.Status != READY_ALTERATION_FAILED {
			return c.JsonBadRequest("Model in invalid status for re-training")
		} else if model.CanTrain == 0 {
			return c.JsonBadRequest("Model can not be trained!")
		}

		c.Logger.Info("Expanding definitions for models", "id", model.Id)

		classesUpdated := false

		failed := func() *Error {
			if classesUpdated {
				ResetClasses(c, model)
			}

			ModelUpdateStatus(c, model.Id, READY_RETRAIN_FAILED)
			return c.E500M("Failed to retrain model", err)
		}

		var def struct {
			Id             string
			TargetAccuracy int `db:"target_accuracy"`
		}

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

		type C struct {
			Id         string
			ClassOrder int `db:"class_order"`
		}

		err = c.StartTx()
		if err != nil {
			return failed()
		}

		classes, err := GetDbMultitple[C](
			c,
			"model_classes where model_id=$1 and status=$2 order by class_order asc",
			model.Id,
			CLASS_STATUS_TO_TRAIN,
		)
		if err != nil {
			_err := c.RollbackTx()
			if _err != nil {
				c.Logger.Error("Two errors happended rollback failed", "err", _err)
			}
			return failed()
		}

		if len(classes) == 0 {
			c.Logger.Error("No classes are available!")
			_err := c.RollbackTx()
			if _err != nil {
				c.Logger.Error("Two errors happended rollback failed", "err", _err)
			}
			return failed()
		}

		//Update the classes
		{
			_, err = c.Exec("update model_classes set status=$1 where status=$2 and model_id=$3", CLASS_STATUS_TRAINING, CLASS_STATUS_TO_TRAIN, model.Id)
			if err != nil {
				_err := c.RollbackTx()
				if _err != nil {
					c.Logger.Error("Two errors happended rollback failed", "err", _err)
				}
				return failed()
			}

			err = c.CommitTx()
			if err != nil {
				_err := c.RollbackTx()
				if _err != nil {
					c.Logger.Error("Two errors happended rollback failed", "err", _err)
				}
				return failed()
			}

			classesUpdated = true
		}

		var newHead = struct {
			DefId      string           `db:"def_id"`
			RangeStart int              `db:"range_start"`
			RangeEnd   int              `db:"range_end"`
			Status     DefinitionStatus `db:"status"`
		}{
			def.Id, classes[0].ClassOrder, classes[len(classes)-1].ClassOrder, DEFINITION_STATUS_INIT,
		}

		_, err = InsertReturnId(c.GetDb(), &newHead, "exp_model_head", "id")
		if err != nil {
			return failed()
		}

		_, err = c.Db.Exec("update models set status=$1 where id=$2;", READY_RETRAIN, model.Id)
		if err != nil {
			return c.E500M("Failed to update model status", err)
		}

		newTask := struct {
			UserId   string   `db:"user_id"`
			ModelId  string   `db:"model_id"`
			TaskType TaskType `db:"task_type"`
			Status   int      `db:"status"`
		}{
			UserId:   c.User.Id,
			ModelId:  model.Id,
			TaskType: TASK_TYPE_RETRAINING,
			Status:   1,
		}

		id, err := InsertReturnId(c, &newTask, "tasks", "id")
		if err != nil {
			return c.E500M("Failed to create task", err)
		}

		return c.SendJSON(JustId{Id: id})
	})

	handle.Get("/model/epoch/update", func(c *Context) *Error {
		f := c.R.URL.Query()

		accuracy := 0.0

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

		accuracy = accuracy * 100

		model_id := f.Get("model_id")
		def_id := f.Get("definition")
		epoch, err := strconv.Atoi(f.Get("epoch"))
		if err != nil {
			return c.JsonBadRequest("Epoch is not a number")
		}

		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)
			return c.JsonBadRequest("Definition not on status 3(training)")
		}

		c.Logger.Debug("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)
		}

		c.ShowMessage = false
		return nil
	})

	handle.Get("/model/head/epoch/update", func(c *Context) *Error {
		f := c.R.URL.Query()

		accuracy := 0.0

		if !CheckId(f, "head_id") || CheckEmpty(f, "epoch") || !CheckFloat64(f, "accuracy", &accuracy) {
			return c.JsonBadRequest("Invalid: model_id or definition or epoch or accuracy")
		}

		accuracy = accuracy * 100

		head_id := f.Get("head_id")
		epoch, err := strconv.Atoi(f.Get("epoch"))
		if err != nil {
			return c.JsonBadRequest("Epoch is not a number")
		}

		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)
			return c.JsonBadRequest("Head not on status 3(training)")
		}

		c.Logger.Debug("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)
		}

		c.ShowMessage = false
		return nil
	})
}