package main

import "C"
import (
	pge "go-pge/pge"
	probs "go-pge/problems"
	rand "math/rand"
	rt "runtime"
	"sort"
	"strings"
	"unsafe"
)

func main() {

}

var cps *pge.PgeSearch
var tp *probs.TreeParams
var epcomm *probs.ExprProblemComm
var trainDatas []*probs.PointSet
var testDatas []*probs.PointSet
var testDatasN int
var trainDatasN int
var ep *probs.ExprProblem
var per_eqns []*probs.ExprReportArray

var trainData *probs.PointSet
var testData *probs.PointSet

var probIndepNames []string
var probDepndNames []string

var stepResult probs.ExprReportArray
var lastStepRes *probs.ExprReport
var resultNr int

var nCurrentCoeff int
var nMaxCoeff int
var sortTypeNr int
var nLastResults int

//export InitSearch
func InitSearch(maxGen int, pgeRptEpoch int, pgeRptCount int, pgeArchiveCap int, peelCnt int, evalrCount int, zeroEpsilon float64, initMethod *C.char, growMethod *C.char, sortType int) {
	cps = new(pge.PgeSearch)

	goInitMethod := C.GoString(initMethod)
	goGrowMethod := C.GoString(growMethod)

	testDatasN = 0
	trainDatasN = 0
	nLastResults = 0
	sortTypeNr = sortType

	cps.CreateDS(maxGen, pgeRptEpoch, pgeRptCount, pgeArchiveCap, peelCnt, evalrCount, zeroEpsilon, goInitMethod, goGrowMethod, sortType)

	cps.SetPeelCount(peelCnt)
	cps.SetGrowMethod(goGrowMethod)
	cps.SetInitMethod(goInitMethod)
	cps.SetEvalrCount(evalrCount)
}

//export InitTreeParams
func InitTreeParams(Roots *C.char, Nodes *C.char, NonTrig *C.char, Leafs *C.char, nUseableVars int, MaxSize int, MinSize int, MaxDepth int, MinDepth int) {
	tp = new(probs.TreeParams)
	goRoots := C.GoString(Roots)
	goNodes := C.GoString(Nodes)
	goNonTrig := C.GoString(NonTrig)
	goLeafs := C.GoString(Leafs)
	goUsableVars := CreateUsableVars(nUseableVars)
	tp.CreateTreeParams(goRoots, goNodes, goNonTrig, goLeafs, goUsableVars, MaxSize, MinSize, MaxDepth, MinDepth)
}

//export AddTrainData
func AddTrainData(indepNames *C.char, depndNames *C.char, matrix unsafe.Pointer, nEntries int) {
	trainDatas = make([]*probs.PointSet, 1)
	trainData = new(probs.PointSet)

	goIndepNames := strings.Split(C.GoString(indepNames), " ")
	goDepndNames := strings.Split(C.GoString(depndNames), " ")

	probIndepNames = goIndepNames
	probDepndNames = goDepndNames

	trainData.InitTrainData(goIndepNames, goDepndNames, matrix, nEntries)

	trainDatas[trainDatasN] = trainData
	trainDatasN++
}

//export AddTestData
func AddTestData(indepNames *C.char, depndNames *C.char, matrix unsafe.Pointer, nEntries int) {
	testDatas = make([]*probs.PointSet, 1)
	testData = new(probs.PointSet)

	goIndepNames := strings.Split(C.GoString(indepNames), " ")
	goDepndNames := strings.Split(C.GoString(depndNames), " ")

	testData.InitTrainData(goIndepNames, goDepndNames, matrix, nEntries)

	testDatas[testDatasN] = testData
	testDatasN++
}

//export InitProblem
func InitProblem(Name *C.char, MaxIter int, HitRatio float64, SearchVar int, ProblemTypeString *C.char, numProcs int) {
	ep = new(probs.ExprProblem)

	ep.IndepNames = probIndepNames
	ep.DepndNames = probDepndNames

	goName := C.GoString(Name)
	goProblemTypeString := C.GoString(ProblemTypeString)

	ep.CreatePS(goName, MaxIter, HitRatio, SearchVar, goProblemTypeString)
	ep.TreeCfg = tp
	ep.UsableVars = tp.UsableVars

	cps.SetMaxIter(MaxIter)
	per_eqns = make([]*probs.ExprReportArray, 1)

	ep.Train = trainDatas
	ep.Test = testDatas

	initDone := make(chan int)

	rt.GOMAXPROCS(numProcs)
	rand.Seed(rand.Int63())

	epcomm = new(probs.ExprProblemComm)

	epcomm.Cmds = make(chan int)
	epcomm.Rpts = make(chan *probs.ExprReportArray, 64)
	epcomm.Res = make(chan *probs.ExprReportArray, 64) //Stops at 64 if nothing gets retrieved
	epcomm.Gen = make(chan [2]int, 64)

	cps.SetProb(ep)
	cps.Init(initDone, ep, "", epcomm)
	//cps.SetSort(sortTypeNr) really? sorttype is hardcoded in PGE!
}

//export GetMaxIterW
func GetMaxIterW() int {
	return cps.GetMaxIter()
}

//export SetMaxIterW
func SetMaxIterW(iter int) {
	cps.SetMaxIter(iter)
}

//export SetPeelCountW
func SetPeelCountW(cnt int) {
	cps.SetPeelCount(cnt)
}

//export SetInitMethodW
func SetInitMethodW(init string) {
	cps.SetInitMethod(init)
}

//export SetGrowMethodW
func SetGrowMethodW(grow string) {
	cps.SetGrowMethod(grow)
}

//export SetEvalrCountW
func SetEvalrCountW(cnt int) {
	cps.SetEvalrCount(cnt)
}

//export EvaluateW
func EvaluateW() {
	cps.Evaluate()
}

//export RunW
func RunW() {
	cps.Run()
}

//export LoopW
func LoopW() {
	cps.Loop()
}

//export GetStepResult
func GetStepResult(testscore *int, ncoeff *int) *C.char {
	*testscore = 0

	if resultNr < len(stepResult) && stepResult != nil {
		elem := stepResult[resultNr]
		resultNr++
		lastStepRes = elem
		if elem != nil {
			*testscore = elem.TestScore()
			*ncoeff = len(elem.Coeff())
			nCurrentCoeff = 0
			nMaxCoeff = len(elem.Coeff())
			return C.CString(elem.Expr().String())
		}
	}
	return C.CString("")
}

//export GetCoeffResult
func GetCoeffResult() float64 {
	var res float64 = 0
	if nCurrentCoeff < nMaxCoeff {
		coeffs := lastStepRes.Coeff()
		res = coeffs[nCurrentCoeff]
	}
	nCurrentCoeff++
	return res
}

func accumExprs(eqns *probs.ExprReportArray) probs.ExprReportArray {
	union := *eqns //make(probs.ExprReportArray, 0)
	//union = append(union, (*eqns)[:]...)

	// remove duplicates
	sort.Sort(union)
	last := 0
	for last < len(union) && union[last] == nil {
		last++
	}
	for i := last + 1; i < len(union); i++ {
		if union[i] == nil {
			continue
		}
		if union[i].Expr().AmIAlmostSame(union[last].Expr()) {
			union[i] = nil
		} else {
			last = i
		}
	}

	queue := probs.NewQueueFromArray(union)
	queue.SetSort(probs.GPSORT_PARETO_TST_ERR)
	queue.Sort()

	ds_eqns := make(probs.ExprReportArray, 20)
	copy(ds_eqns, union[:20])

	return ds_eqns
}

//export GetBestResults
func GetBestResults() int {
	cps.ReportExpr(true)

	rpt := <-epcomm.Rpts
	stepResult = accumExprs(rpt)
	resultNr = 0

	return len(stepResult)
}

//export StepW
func StepW() int {
	resultNr = 0
	stepResult = nil
	nRes := cps.Step()

	resu, ok := <-epcomm.Res
	resultNr = nLastResults

	nNew := 0
	if ok && resu != nil {
		stepResult = *resu
		nNew = nRes
		nLastResults = 0
	}

	return nNew
}

//export Clear
func Clear() {
	cps.Stop()
}

//export CheckStop
func CheckStop() bool {
	if cps.GetIter() > GetMaxIterW() {
		return true
	}
	return false
}

func CreateUsableVars(nEntries int) []int {
	res := make([]int, nEntries)
	for i := 0; i < nEntries; i++ {
		res[i] = i
	}
	return res
}

func CIntArrayToSlice(matrix unsafe.Pointer, nEntries int) []int {
	res := make([]int, nEntries)
	for i := 0; i < nEntries; i++ {
		item := (*int)(unsafe.Pointer(uintptr(unsafe.Pointer(matrix)) + unsafe.Sizeof(int(0))*uintptr(i)))
		res[i] = *item

	}
	return res
}

//export CleanW
func CleanW() {
	cps.Clean()
}